Update app.py

app.py

@@ -5,16 +5,23 @@ import matplotlib.pyplot as plt
 from io import BytesIO
 from llama_index import VectorStoreIndex, SimpleDirectoryReader
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.llms.huggingface import HuggingFaceLLM
 import dotenv
 import re
+import requests
 
 # Load environment variables
 dotenv.load_dotenv()
 
-# Configure Hugging Face models
+# Configure Hugging Face API
+API_URL = "https://api-inference.huggingface.co/models/sarvamai/sarvam-2b-v0.5"
+headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}
+
+def query_huggingface_api(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()
+
+# Configure embedding model
 embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
-llm_model = HuggingFaceLLM(model_name="sarvamai/sarvam-2b-v0.5", api_token=os.getenv("HUGGINGFACE_API_KEY"))
 
 def write_to_file(content, filename="./files/test.pdf"):
     os.makedirs(os.path.dirname(filename), exist_ok=True)
@@ -22,9 +29,6 @@ def write_to_file(content, filename="./files/test.pdf"):
         f.write(content)
 
 def extract_financial_data(document_text):
-    """
-    Extracts financial data such as revenue and dates from the document text.
-    """
     financial_data = {
         "Revenue": [],
         "Date": []
@@ -35,7 +39,7 @@ def extract_financial_data(document_text):
 
     for i, line in enumerate(lines):
         if any(keyword in line.lower() for keyword in ["revenue", "total revenue", "sales"]):
-            for j in range(i + 1, i + 6):  # Look ahead a few lines for numbers
+            for j in range(i + 1, i + 6):
                 matches = revenue_pattern.findall(lines[j])
                 if matches:
                     for match in matches:
@@ -64,26 +68,28 @@ def load_data(documents):
     return index
 
 def generate_summary(index, document_text, query):
-    query_engine = index.as_query_engine(llm_model=llm_model)
-    response = query_engine.query(f"""
-    You are a financial analyst. Your task is to provide a comprehensive analysis of the financial document.
-    Analyze the following document and respond to the query:
-    {document_text}
-    
-    Query: {query}
-    
-    If the query is too general, respond with:
-    Please cover the following aspects:
-    1. Revenue and profit trends
-    2. Key financial metrics
-    3. Major financial events and decisions
-    4. Comparison with previous periods
-    5. Future outlook or forecasts
-    6. Any notable financial risks or opportunities
-    
-    Provide a clear, concise, and professional response.
-    """)
-    return response.response
+    query_engine = index.as_query_engine()
+    llm_response = query_huggingface_api({
+        "inputs": f"""
+        You are a financial analyst. Your task is to provide a comprehensive analysis of the financial document.
+        Analyze the following document and respond to the query:
+        {document_text}
+        
+        Query: {query}
+        
+        If the query is too general, respond with:
+        Please cover the following aspects:
+        1. Revenue and profit trends
+        2. Key financial metrics
+        3. Major financial events and decisions
+        4. Comparison with previous periods
+        5. Future outlook or forecasts
+        6. Any notable financial risks or opportunities
+        
+        Provide a clear, concise, and professional response.
+        """
+    })
+    return llm_response.get("generated_text", "No response from model.")
 
 def generate_comparison_graph(data):
     if not data["Date"] or not data["Revenue"]:
@@ -145,4 +151,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-