Update langchain_movie_search.py

langchain_movie_search.py

@@ -1,26 +1,18 @@
 import os
 from typing import List
-import argparse
 from dotenv import load_dotenv
 import pymongo
-import certifi
 from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
 from langchain.chains import create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain_core.documents import Document
 from langchain_core.prompts import PromptTemplate
 import gradio as gr
 from gradio.themes.base import Base
-from flask import Flask
 
 __author__ = "Chirag Kamble"
 
 
-# Flask App
-# app = Flask(__name__)
-
-
 class MoviesSearch:
     """
     Class to perform Vector Index Search using MongoDB and LLM search using Langchain on Movies
@@ -44,13 +36,10 @@ class MoviesSearch:
         self.client: pymongo.synchronous.mongo_client.MongoClient = pymongo.MongoClient(mongodb_connection_url,
                                                                                         serverSelectionTimeoutMS=60000,
                                                                                         tls=True,
-                                                                                        # tlsCAFile=certifi.where(),
                                                                                         connect=False,
                                                                                         tlsAllowInvalidCertificates=True,
                                                                                         directConnection=False,
-                                                                                        # tlsInsecure=True,
                                                                                         maxPoolSize=100,
-                                                                                        # minPoolSize=0,
                                                                                         maxIdleTimeMS=60000,
                                                                                         waitQueueTimeoutMS=60000,
                                                                                         connectTimeoutMS=60000,
@@ -80,20 +69,15 @@ class MoviesSearch:
         Generate vector embeddings
         """
         new_doc_list: List[Document] = []
-        for doc in self.sample_movies_collection.find({"fullplot": {"$exists": True}}).limit(1000):
-            new_doc: Document = Document(
-                page_content=doc["fullplot"],
-                metadata={"source": "Collection sample_mflix",
-                          "movie-title": doc["title"],
-                          "movie-plot": doc["fullplot"],
-                          "text": doc["fullplot"]}
-            )
-            new_doc_list.append(new_doc)
-        self.retrieve_vector_store.from_documents(
-            documents=new_doc_list,
-            embedding=self.hf_plot_embedding,
-            collection=self.langchain_movies_collection
-        )
+        for doc in self.sample_movies_collection.find({"fullplot": {"$exists": True}}).limit(9000):
+            new_doc_list.append({
+                "movie-title": doc["title"],
+                "movie-plot": doc["fullplot"],
+                "text": doc["fullplot"],
+                "embedding": self.hf_plot_embedding.embed_query(doc["fullplot"])
+            })
+
+        self.langchain_movies_collection.insert_many(new_doc_list)
 
     def query_data(self, query: str):
         """
@@ -111,7 +95,6 @@ class MoviesSearch:
             return_full_text=True,
         )
 
-        # retriever = self.retrieve_vector_store.as_retriever()
         retriever = self.retrieve_vector_store.as_retriever()
         
         prompt = PromptTemplate.from_template(template="{context}", template_format="f-string")
@@ -142,36 +125,11 @@ class MoviesSearch:
         self.client.close()
 
 
-# @app.route("/", methods=["GET"])
 def gradio_interface(cmd=None):
     movie_search = MoviesSearch()
     # movie_search.generate_insert_embeddings()
     movie_search.run_website()
 
-    # if cmd == "generate_embeddings":
-    #     movie_search.generate_insert_embeddings()
-    # elif cmd == "run":
-    #     movie_search.run_website()
-
 
 if __name__ == "__main__":
-    # Create the parser
-    # parser = argparse.ArgumentParser(description='Script to suggest movies based on user description/query')
-    #
-    # # Add arguments
-    # parser.add_argument("-g", "--generate_embeddings", action="store_true", help="Generate/Re-generate Embeddings")
-    # parser.add_argument("-r", "--run", action="store_true", help="Age of the person")
-    #
-    # # Parse arguments
-    # args = parser.parse_args()
-    #
-    # if args.generate_embeddings:
-    #     gradio_interface(cmd="generate_embeddings")
-    # elif args.run:
-    #     gradio_interface(cmd="run")
-
-    # app.run(host="0.0.0.0", port=os.getenv("PORT", 5000), debug=True)
-    # app.run(host="0.0.0.0", debug=True)
-    # app.run(debug=True)
-
     gradio_interface()