Update app.py

app.py

@@ -29,6 +29,8 @@ load_dotenv()
 
 os.environ['API_KEY_PROJ3'] = os.getenv('API_KEY_PROJ3')
 
+collection_name = 'collection'
+
 client = OpenAI(
     base_url="https://api.endpoints.anyscale.com/v1",
     api_key=os.environ['API_KEY_PROJ3']
@@ -38,15 +40,20 @@ client = OpenAI(
 embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
 
 # Load the persisted vectorDB
+vectorstore_persisted = Chroma(
+    collection_name=collection_name,
+    persist_directory='./proj3_db',
+    embedding_function=embedding_model
+)
 persisted_vectordb_location = './proj3_db'
 
 # Prepare the logging functionality
 
 log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
 log_folder = log_file.parent
- 
+
 scheduler = CommitScheduler(
-    repo_id="dataset",
+    repo_id="---------",
     repo_type="dataset",
     folder_path=log_folder,
     path_in_repo="data",
@@ -76,3 +83,75 @@ Here are some documents that are relevant to the question mentioned below.
 ###Question
 {question}
 """
+
+# Define the predict function that runs when 'Submit' is clicked or when a API request is made
+def predict(user_input,company):
+
+    filter = "dataset/"+company+"-10-k-2023.pdf"
+    relevant_document_chunks = vectorstore_persisted.similarity_search(user_input, k=5, filter={"source":filter})
+
+    # Create context_for_query
+    context_for_query = ". ".join(relevant_document_chunks)
+
+    # Create messages
+    prompt = [
+        {'role': 'system', 'content': qna_system_message},
+        {'role': 'user', 'content': qna_user_message_template.format(
+            context=context_for_query,
+            question=user_input
+        )
+         }
+    ]
+    model_name = 'mlabonne/NeuralHermes-2.5-Mistral-7B'
+    # Get response from the LLM
+    try:
+        response = client.chat.completions.create(
+            model=model_name,
+            messages=prompt,
+            temperature=0
+        )
+
+        prediction = response.choices[0].message.content.strip()
+    except Exception as e:
+        prediction = f'Sorry, I encountered the following error: \n {e}'
+
+    # 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
+
+    with scheduler.lock:
+        with log_file.open("a") as f:
+            f.write(json.dumps(
+                {
+                    'user_input': user_input,
+                    'retrieved_context': context_for_query,
+                    'model_response': prediction
+                }
+            ))
+            f.write("\n")
+
+    return prediction
+
+# Set-up the Gradio UI
+# Add text box and radio button to the interface
+# The radio button is used to select the company 10k report in which the context needs to be retrieved.
+lst_companies = ['aws', 'google', 'IBM', 'Meta', 'msft']
+textbox = gr.Textbox('Input user')
+company = gr.Radio('Company', lst_companies)
+
+model_output = gr.Label(label="Charge predictor")
+
+# Create the interface
+# For the inputs parameter of Interface provide [textbox,company]
+demo = gr.Interface(
+    fn=predict,
+    inputs=[textbox, company],
+    outputs=model_output,
+    title="Charge Predictor",
+    description="This API allows you to predict the charge of insurace",
+    allow_flagging="auto",
+    concurrency_limit=8
+)
+
+demo.queue()
+demo.launch(share=False)