Update app.py

app.py

@@ -1,33 +1,26 @@
-
-# Import the necessary Libraries
 import os
 import uuid
 import json
 
 import gradio as gr
-#import openai
-#import load_dotenv
-!pip install openai
 from openai import OpenAI
-
 from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
 from langchain_community.vectorstores import Chroma
-
 from huggingface_hub import CommitScheduler
 from pathlib import Path
 from dotenv import load_dotenv
 
-
-# Create Client
+# Load environment variables from .env file
 load_dotenv()
 
-os.environ["OPENAI_API_KEY"] = "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"
+# Set OpenAI API key
+openai.api_key = os.getenv("OPENAI_API_KEY")  # Make sure OPENAI_API_KEY is in your .env file
 
+# Initialize OpenAI client
 client = OpenAI()
 
+# Set up embeddings and vectorstore
 embeddings = SentenceTransformerEmbeddings(model_name="thenlper/gte-large")
-# Define the embedding model and the vectorstore
-
 collection_name = 'report-10k-2024'
 
 vectorstore_persisted = Chroma(
@@ -36,16 +29,12 @@ vectorstore_persisted = Chroma(
     embedding_function=embeddings
 )
 
-# Load the persisted vectorDB
-
 retriever = vectorstore_persisted.as_retriever(
     search_type='similarity',
     search_kwargs={'k': 5}
 )
 
-
-# Prepare the logging functionality
-
+# Set up logging
 log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
 log_folder = log_file.parent
 
@@ -58,13 +47,10 @@ scheduler = CommitScheduler(
 )
 
 # Define the Q&A system message
-
 qna_system_message = """
 You are an AI assistant to help Finsights Grey Inc., an innovative financial technology firm, develop a Retrieval-Augmented Generation (RAG) system to automate the extraction, summarization, and analysis of information from 10-K reports. Your knowledge base was last updated in August 2023.
-
 User input will have the context required by you to answer user questions. This context will begin with the token: ###Context.
 The context contains references to specific portions of a 10-K report relevant to the user query.
-
 User questions will begin with the token: ###Question.
 Your response should only be about the question asked and the context provided.
 Answer only using the context provided.
@@ -72,44 +58,29 @@ Do not mention anything about the context in your final answer.
 If the answer is not found in the context, it is very important for you to respond with "I don't know."
 Always quote the source when you use the context. Cite the relevant source at the end of your response under the section - Source:
 Do not make up sources. Use the links provided in the sources section of the context and nothing else. You are prohibited from providing other links/sources.
-Here is an example of how to structure your response:
-
-Answer:
-[Answer]
-
-Source:
-[Source]
 """
 
-# Define the user message template
 qna_user_message_template = """
 ###Context
 Here are some documents that are relevant to the question.
 {context}
-```
 {question}
-```
-"""
 
-# Define the predict function that runs when 'Submit' is clicked or when a API request is made
-def predict(user_input,company):
+css
+Copy code
+"""
 
-    filter = "dataset/"+company+"-10-k-2023.pdf"
-    relevant_document_chunks = vectorstore_persisted.similarity_search(user_input, k=5, filter={"source":filter})
+# Define the predict function
+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
+    # Create context for query
     context_list = [d.page_content for d in relevant_document_chunks]
     context_for_query = ".".join(context_list)
 
     # Create messages
-    prompt = [
-        {'role':'system', 'content': qna_system_message},
-        {'role': 'user', 'content': qna_user_message_template.format(
-            context=context_for_query,
-            question=user_input
-            )
-        }
-    ]
+    prompt = [        {'role': 'system', 'content': qna_system_message},        {'role': 'user', 'content': qna_user_message_template.format(            context=context_for_query,            question=user_input        )}    ]
 
     # Get response from the LLM
     try:
@@ -122,12 +93,9 @@ def predict(user_input,company):
         prediction = response.choices[0].message.content
 
     except Exception as e:
-        prediction = 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
+        prediction = str(e)
 
+    # Log inputs and outputs to a local log file
     with scheduler.lock:
         with log_file.open("a") as f:
             f.write(json.dumps(
@@ -144,40 +112,26 @@ def predict(user_input,company):
 
 def get_predict(question, company):
     # Implement your prediction logic here
-    if company == "AWS":
-        # Perform prediction for AWS
-        selectedCompany = "aws"
-    elif company == "IBM":
-        # Perform prediction for IBM
-        selectedCompany = "IBM"
-    elif company == "Google":
-        # Perform prediction for Google
-       selectedCompany = "Google"
-    elif company == "Meta":
-        # Perform prediction for Meta
-        selectedCompany = "meta"
-    elif company == "Microsoft":
-        # Perform prediction for Microsoft
-        selectedCompany = "msft"
-    else:
+    company_map = {
+        "AWS": "aws",
+        "IBM": "IBM",
+        "Google": "Google",
+        "Meta": "meta",
+        "Microsoft": "msft"
+    }
+    
+    selected_company = company_map.get(company)
+    if not selected_company:
         return "Invalid company selected"
-
-    output = predict(question, selectedCompany)
-    return output
+    
+    return predict(question, selected_company)
 
 # 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.
-
-# Create the interface
-# For the inputs parameter of Interface provide [textbox,company]
-
 with gr.Blocks(theme="gradio/seafoam@>=0.0.1,<0.1.0") as demo:
     with gr.Row():
         company = gr.Radio(["AWS", "IBM", "Google", "Meta", "Microsoft"], label="Select a company")
         question = gr.Textbox(label="Enter your question")
 
-
     submit = gr.Button("Submit")
     output = gr.Textbox(label="Output")
 
@@ -188,4 +142,4 @@ with gr.Blocks(theme="gradio/seafoam@>=0.0.1,<0.1.0") as demo:
     )
 
 demo.queue()
-demo.launch()
+demo.launch()