Upload app.py

app.py

@@ -0,0 +1,120 @@
+
+import gradio as gr
+from sentence_transformers import SentenceTransformer
+import chromadb
+import pandas as pd
+import os
+
+# Load the sentence transformer model
+model = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Initialize the ChromaDB client
+client = chromadb.Client()
+
+# Function to build the database from CSV
+def build_database():
+    # Read the CSV file
+    df = pd.read_csv('collection_data.csv')
+    
+    # Create a collection
+    collection_name = 'Dataset-10k-companies'
+    
+    # Delete the existing collection if it exists
+    if collection_name in client.list_collections():
+        client.delete_collection(name=collection_name)
+    
+    # Create a new collection
+    collection = client.create_collection(name=collection_name)
+    
+    # Add the data from the DataFrame to the collection
+    collection.add(
+        documents=df['documents'].tolist(),
+        ids=df['ids'].tolist(),
+        metadatas=df['metadatas'].apply(eval).tolist(),
+        embeddings=df['embeddings'].apply(lambda x: eval(x.replace(',,', ','))).tolist()
+    )
+    
+    return collection
+
+# Build the database when the app starts
+collection = build_database()
+
+# Function to get relevant chunks
+def get_relevant_chunks(query, collection, top_n=3):
+    query_embedding = model.encode(query).tolist()
+    results = collection.query(query_embeddings=[query_embedding], n_results=top_n)
+    
+    relevant_chunks = []
+    for i in range(len(results['documents'][0])):
+        chunk = results['documents'][0][i]
+        source = results['metadatas'][0][i]['source']
+        page = results['metadatas'][0][i]['page']
+        relevant_chunks.append((chunk, source, page))
+    
+    return relevant_chunks
+
+# Function to get LLM response
+def get_llm_response(prompt, max_attempts=3):
+    full_response = ""
+    for attempt in range(max_attempts):
+        try:
+            response = client.complete(prompt, max_tokens=1000)  # Increase max_tokens if possible
+            chunk = response.text.strip()
+            full_response += chunk
+            if chunk.endswith((".", "!", "?")):  # Check if response seems complete
+                break
+            else:
+                prompt = "Please continue from where you left off:\n" + chunk[-100:]  # Use the last 100 chars as context
+        except Exception as e:
+            print(f"Attempt {attempt + 1} failed with error: {e}")
+    return full_response
+
+# Prediction function
+def predict(company, user_query):
+    # Modify the query to include the company name
+    modified_query = f"{user_query} for {company}"
+    
+    # Get relevant chunks
+    relevant_chunks = get_relevant_chunks(modified_query, collection)
+    
+    # Prepare the context string
+    context = ""
+    for chunk, source, page in relevant_chunks:
+        context += chunk + "\n"
+        context += f"[Source: {source}, Page: {page}]\n\n"
+    
+    # Generate answer
+    prompt = f"Based on the following context, answer the question: {modified_query}\n\nContext:\n{context}"
+    answer = get_llm_response(prompt)
+        # 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 answer
+
+# Create Gradio interface
+company_list = ["MSFT", "AWS", "Meta", "Google", "IBM"]
+iface = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Radio(company_list, label="Select Company"),
+        gr.Textbox(lines=2, placeholder="Enter your query here...", label="User Query")
+    ],
+    outputs=gr.Textbox(label="Generated Answer"),
+    title="Company Reports Q&A",
+    description="Query the vector database and get an LLM response based on the documents in the collection."
+)
+
+# Launch the interface
+iface.launch()