Create app.py

app.py

@@ -0,0 +1,185 @@
+import gradio as gr
+from sentence_transformers import SentenceTransformer
+import chromadb
+import pandas as pd
+import os
+import json
+from pathlib import Path
+from llama_index.llms.anyscale import Anyscale
+from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
+
+# Load the sentence transformer model with 384 dimensions
+model = SentenceTransformer('all-MiniLM-L6-v2')
+# model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')  # Ensure this model outputs 384-dimensional embeddings
+
+# Check the dimensionality of the model
+embedding_dim = model.get_sentence_embedding_dimension()
+print(f"Embedding Dimension: {embedding_dim}")  # This should print 384
+
+# Initialize the ChromaDB client
+chroma_client = chromadb.Client()
+
+# Function to build the database from CSV
+def build_database():
+    # Read the CSV file
+    df = pd.read_csv('vector_store.csv')
+
+    # Create a collection
+    collection_name = 'Dataset-10k-companies'
+    
+    # Delete the existing collection if it exists
+    # chroma_client.delete_collection(name=collection_name)
+    
+    # Create a new collection
+    collection = chroma_client.create_collection(name=collection_name)
+    
+    # Ensure embeddings are 384-dimensional
+    def ensure_dimensionality(embedding):
+        embedding = eval(embedding.replace(',,', ','))
+        if len(embedding) != 384:
+            raise ValueError("Embedding dimensionality is incorrect")
+        return embedding
+    
+    # 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(ensure_dimensionality).tolist()
+    )
+    
+    return collection
+
+# Build the database when the app starts
+collection = build_database()
+
+# Access the Anyscale API key from environment variables
+anyscale_api_key = os.environ.get('anyscale_api_key')
+
+# Instantiate the Anyscale client
+client = Anyscale(api_key=anyscale_api_key, model="meta-llama/Llama-2-70b-chat-hf")
+
+# 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
+
+# Define system message for LLM
+qna_system_message = """
+You are an assistant to Finsights analysts. Your task is to provide relevant information about the financial performance of the companies followed by Finsights.
+User input will include the necessary context for you to answer their questions. This context will begin with the token: ###Context.
+The context contains references to specific portions of documents relevant to the user's query, along with source links.
+The source for a context will begin with the token: ###Source.
+When crafting your response:
+1. Select only the context relevant to answer the question.
+2. Include the source links in your response.
+3. User questions will begin with the token: ###Question.
+4. If the question is irrelevant to Finsights, respond with: "I am an assistant for Finsight Docs. I can only help you with questions related to Finsights."
+Adhere to the following guidelines:
+- Your response should only address the question asked and nothing else.
+- Answer only using the context provided.
+- Do not mention anything about the context in your final answer.
+- If the answer is not found in the context, 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 only the links provided in the sources section of the context. You are prohibited from providing other links/sources.
+Here is an example of how to structure your response:
+Answer:
+[Answer]
+Source:
+[Source]
+"""
+# Create a user message template
+qna_user_message_template = """
+###Context
+Here are some documents and their source links that are relevant to the question mentioned below.
+{context}
+###Question
+{question}
+"""
+
+# 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):
+    try:
+        # 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"
+
+        # Prepare the user message
+        user_message = qna_user_message_template.format(context=context, question=user_query)
+
+        # Craft the prompt to pass to the Llama model
+        prompt = f"{qna_system_message}\n\n{qna_user_message_template.format(context=context, question=user_query)}"
+
+        # Generate the response using the Llama model through Anyscale
+        answer = get_llm_response(prompt)
+
+        # Extract the generated response
+        # answer = response.text.strip()
+        
+        # Log the interaction
+        log_interaction(company, user_query, context, answer)
+        
+        return answer
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+# Function to log interactions
+def log_interaction(company, user_query, context, answer):
+    log_file = Path("interaction_log.jsonl")
+    with log_file.open("a") as f:
+        json.dump({
+            'company': company,
+            'user_query': user_query,
+            'context': context,
+            'answer': answer
+        }, f)
+        f.write("\n")
+
+# 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(share=True)