Update app.py

app.py

@@ -2,17 +2,15 @@ import gradio as gr
 import pandas as pd
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
-import pymongo
-import os
-from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
 
 # Load a smaller portion of the dataset
-dataset = load_dataset("MongoDB/embedded_movies", split='train[:10%]')
+dataset = load_dataset("MongoDB/embedded_movies", split='train[:5%]')
 dataset_df = pd.DataFrame(dataset)
 
 # Data cleaning and preprocessing
-dataset_df = dataset_df.dropna(subset=["fullplot"])
+dataset_df = dataset_df.dropna(subset=["fullplot"]).reset_index(drop=True)
 dataset_df = dataset_df.drop(columns=["plot_embedding"])
 
 # Load a smaller embedding model
@@ -25,52 +23,31 @@ def get_embedding(text: str) -> list:
     embedding = embedding_model.encode(text)
     return embedding.tolist()
 
-# Process embeddings in batches
-batch_size = 100
+# Generate embeddings for all plots
+all_embeddings = []
+batch_size = 32
 for i in range(0, len(dataset_df), batch_size):
-    batch = dataset_df.iloc[i:i+batch_size]
-    batch["embedding"] = batch["fullplot"].apply(get_embedding)
-    dataset_df.iloc[i:i+batch_size] = batch
-
-def get_mongo_client(mongo_uri):
-    try:
-        client = pymongo.MongoClient(mongo_uri)
-        print("Connection to MongoDB successful")
-        return client
-    except pymongo.errors.ConnectionFailure as e:
-        print(f"Connection failed: {e}")
-        return None
-
-mongo_uri = os.getenv("MONGO_URI")
-if not mongo_uri:
-    print("MONGO_URI not set in environment variables")
-
-mongo_client = get_mongo_client(mongo_uri)
-db = mongo_client["movies"]
-collection = db["movie_collection_2"]
+    batch = dataset_df['fullplot'].iloc[i:i+batch_size].tolist()
+    batch_embeddings = embedding_model.encode(batch)
+    all_embeddings.extend(batch_embeddings)
 
-# Clear the collection and insert new data in bulk
-collection.delete_many({})
-documents = dataset_df.to_dict("records")
-collection.insert_many(documents)
-print("Data ingestion into MongoDB completed")
+# Add embeddings to the DataFrame
+dataset_df['embedding'] = all_embeddings
 
-# Load all embeddings into memory for faster similarity search
-all_embeddings = np.array(dataset_df["embedding"].tolist())
-all_titles = dataset_df["title"].tolist()
+print("Embeddings generated and added to DataFrame")
 
 def vector_search(user_query):
     query_embedding = get_embedding(user_query)
     if not query_embedding:
         return "Invalid query or embedding generation failed."
 
-    similarities = cosine_similarity([query_embedding], all_embeddings)[0]
+    similarities = cosine_similarity([query_embedding], list(dataset_df['embedding']))[0]
     top_indices = similarities.argsort()[-3:][::-1]
     
     results = []
     for idx in top_indices:
         results.append({
-            "title": all_titles[idx],
+            "title": dataset_df.iloc[idx]["title"],
             "fullplot": dataset_df.iloc[idx]["fullplot"],
             "genres": dataset_df.iloc[idx]["genres"],
             "score": similarities[idx]