Update app.py

app.py

@@ -2,12 +2,13 @@ import gradio as gr
 import pandas as pd
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import pymongo
 import os
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
 
 # Load a smaller portion of the dataset
-dataset = load_dataset("MongoDB/embedded_movies", split='train[:70%]')
+dataset = load_dataset("MongoDB/embedded_movies", split='train[:10%]')
 dataset_df = pd.DataFrame(dataset)
 
 # Data cleaning and preprocessing
@@ -54,73 +55,40 @@ documents = dataset_df.to_dict("records")
 collection.insert_many(documents)
 print("Data ingestion into MongoDB completed")
 
-def vector_search(user_query, collection):
+# Load all embeddings into memory for faster similarity search
+all_embeddings = np.array(dataset_df["embedding"].tolist())
+all_titles = dataset_df["title"].tolist()
+
+def vector_search(user_query):
     query_embedding = get_embedding(user_query)
-    if query_embedding is None:
+    if not query_embedding:
         return "Invalid query or embedding generation failed."
 
-    pipeline = [
-        {
-            "$vectorSearch": {
-                "index": "vector_index",
-                "queryVector": query_embedding,
-                "path": "embedding",
-                "numCandidates": 100,
-                "limit": 3,
-            }
-        },
-        {
-            "$project": {
-                "_id": 0,
-                "fullplot": 1,
-                "title": 1,
-                "genres": 1,
-                "score": {"$meta": "vectorSearchScore"},
-            }
-        },
-    ]
-    results = collection.aggregate(pipeline)
-    return list(results)
+    similarities = cosine_similarity([query_embedding], all_embeddings)[0]
+    top_indices = similarities.argsort()[-3:][::-1]
+    
+    results = []
+    for idx in top_indices:
+        results.append({
+            "title": all_titles[idx],
+            "fullplot": dataset_df.iloc[idx]["fullplot"],
+            "genres": dataset_df.iloc[idx]["genres"],
+            "score": similarities[idx]
+        })
+    
+    return results
 
 def get_search_result(query):
-    get_knowledge = vector_search(query, collection)
+    get_knowledge = vector_search(query)
     search_result = ""
     for result in get_knowledge:
         search_result += f"Title: {result.get('title', 'N/A')}\nGenres: {', '.join(result.get('genres', ['N/A']))}\nPlot: {result.get('fullplot', 'N/A')[:150]}...\n\n"
     return search_result
 
-# Lazy loading of the language model
-model = None
-tokenizer = None
-
-def load_language_model():
-    global model, tokenizer
-    if model is None or tokenizer is None:
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
-        model = AutoModelForCausalLM.from_pretrained("gpt2")
-
 def generate_response(query):
-    load_language_model()
     source_information = get_search_result(query)
-    combined_information = f"Answer the question '{query}' based on these movie details:\n\n{source_information}"
-
-    max_length = tokenizer.model_max_length
-    input_ids = tokenizer(combined_information, return_tensors="pt", max_length=max_length, truncation=True)
-
-    try:
-        response = model.generate(
-            **input_ids,
-            max_new_tokens=100,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,
-            top_k=40,
-            top_p=0.9,
-            temperature=0.7,
-            do_sample=True
-        )
-        return tokenizer.decode(response[0], skip_special_tokens=True)
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
+    response = f"Based on your query '{query}', here are some movie recommendations:\n\n{source_information}\nThese movies match your query based on their plot summaries and genres. Let me know if you'd like more information about any of them!"
+    return response
 
 def query_movie_db(user_query):
     return generate_response(user_query)
@@ -128,7 +96,7 @@ def query_movie_db(user_query):
 description_and_article = """
 Ask this bot to recommend you a movie.
 Checkout [my github repo](https://github.com/kanad13/Movie-Recommendation-Bot) to look at the code that powers this bot.
-Note that the bot truncates replies due to token limitations in the free tier of Hugging Face resources.
+Note that the bot provides concise recommendations based on a limited dataset to ensure optimal performance.
 """
 
 iface = gr.Interface(