Upload app.py

app.py

@@ -34,8 +34,8 @@ client = OpenAI(
 embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
 # Load the persisted vectorDB
 collection_name = 'Dataset-10k'
-persisted_vectordb_location = "./bubuuunel/RAG10K"
-reportdb = Chroma(
+
+dataset_db = Chroma(
     collection_name=collection_name,
     persist_directory='./dataset_db',
     embedding_function=embedding_model
@@ -91,10 +91,9 @@ def predict(user_input,company):
     }
 
     filter = "dataset/"+company+"-10-k-2023.pdf"
-    retreiver = vectorstore_persisted.similarity_search(user_input, k=5, filter={"source":filter})
-
+   
     # Create context_for_query
-    relevant_document_chunks = retriever.get_relevant_documents(user_question)
+    relevant_document_chunks = dataset_db.similarity_search(user_question, k=5, filter = {"source":"dataset/google-10-k-2023.pdf"})
     context_list = [d.page_content for d in relevant_document_chunks]
     context_for_query = ". ".join(context_list)
 
@@ -106,7 +105,7 @@ def predict(user_input,company):
             )
         }
     ]
-   
+
     # Create messages
     try:
       response = client.chat.completions.create(
@@ -117,14 +116,14 @@ def predict(user_input,company):
 
       prediction = response.choices[0].message.content.strip()
     except Exception as e:
-    
+
       prediction = f'Sorry, I encountered the following error: \n {e}'
-    
+
 
 
     # Get response from the LLM
     prediction = response.choices[0].message.content.strip()
-    
+
     # While the prediction is made, log both the inputs and outputs to a local log file
     # While writing to the log file, ensure that the commit scheduler is locked to avoid parallel
     # access
@@ -145,9 +144,9 @@ def predict(user_input,company):
 # Set-up the Gradio UI
 user_input = gr.Textbox (label = 'Query')
 company_input = gr.Radio(
-    ['aws','google','IBM','Meta','msft'], 
+    ['aws','google','IBM','Meta','msft'],
     label = 'company'
-    ) 
+    )
 
 model_output = gr.Textbox (label = 'Response')
 
@@ -162,7 +161,7 @@ model_output = gr.Textbox (label = 'Response')
 demo = gr.Interface(
     fn=predict,
     inputs=[user_input,company_input],
-    outputs=model_output,
+    outputs=prediction,
     title="RAG on 10k-reports",
     description="This API allows you to query on annaul reports",
     concurrency_limit=16