Update langchain_movie_search.py

langchain_movie_search.py

@@ -32,9 +32,12 @@ class MoviesSearch:
         mongodb_connection_url: str = os.getenv("MONGODB_CONNECTION_URL")
         mongodb_db_name: str = os.getenv("MONGODB_DB_NAME")
         mongodb_collection_name: str = os.getenv("MONGODB_COLLECTION_NAME")
+        mongodb_index_name: str = os.getenv("MONGODB_INDEX_NAME")
+        text_key: str = os.getenv("TEXT_KEY")
+        embedding_key: str = os.getenv("EMBEDDING_KEY")
         self.huggingface_api_token: str = os.getenv("HF_TOKEN")
         self.huggingface_text_generation_model: str = os.getenv("HUGGINGFACE_TEXT_GENERATION_MODEL")
-
+        
         # Setup MongoDB connection
         self.client: pymongo.synchronous.mongo_client.MongoClient = pymongo.MongoClient(mongodb_connection_url)
         db: str = mongodb_db_name
@@ -47,26 +50,11 @@ class MoviesSearch:
 
         self.retrieve_vector_store = MongoDBAtlasVectorSearch(collection=self.langchain_movies_collection,
                                                               embedding=self.hf_plot_embedding,
-                                                              embedding_key="embedding",
-                                                              index_name="movies_data_12k_vector_index",
-                                                              text_key="uuid_plot",
+                                                              embedding_key=embedding_key,
+                                                              index_name=mongodb_index_name,
+                                                              text_key=text_key,
                                                               )
 
-    def generate_insert_embeddings(self):
-        """
-        Generate vector embeddings
-        """
-        new_doc_list: List = []
-        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):
         """
         Query data from Atlas Vector Search
@@ -92,7 +80,6 @@ class MoviesSearch:
         formatted_prompt = prompt.format(context=query)
         llm_answer = hf_llm.invoke(formatted_prompt)
         llm_answer = llm_answer.split("\n", 1)[1]
-        print(llm_answer)
 
         return llm_answer
 
@@ -102,12 +89,12 @@ class MoviesSearch:
             gr.Markdown("# Generate Movie Plot using Vector Search + RAG")
             with gr.Row():
                 textbox = gr.Textbox(label="Enter your prompt here:", lines=1,
-                                     placeholder="e.g. Generate a movie of a couple discovering love in war")
+                                     placeholder="e.g. Generate a movie where a couple discovers love during a war")
             with gr.Row():
                 button = gr.Button("Generate")
             with gr.Column():
                 output = gr.Textbox(interactive=False,
-                                    label="Here is a Movie Plot for you. Don't forget to invite us to the premier!",
+                                    label="Here is a Movie Plot for you. Don't forget to invite us to the premier! :)",
                                     autoscroll=False,
                                     show_label=True,
                                     show_copy_button=True,
@@ -115,7 +102,7 @@ class MoviesSearch:
 
             button.click(fn=self.query_data, inputs=textbox, outputs=[output])
 
-        dashboard.launch(debug=True)
+        dashboard.launch(share=True)
 
 
 if __name__ == "__main__":