Update app.py

app.py

@@ -9,13 +9,12 @@ from transformers import (
 import torch
 import pandas as pd
 import json
-from huggingface_hub import login
 
 class FinancialAnalyzer:
     def __init__(self):
         print("Loading models...")
         try:
-            # Initialize TinyLlama with correct path
+            # Initialize TinyLlama
             self.tiny_tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
             self.tiny_model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
             
@@ -27,7 +26,7 @@ class FinancialAnalyzer:
             self.t5_tokenizer = T5Tokenizer.from_pretrained("t5-small")
             self.t5_model = T5ForConditionalGeneration.from_pretrained("t5-small")
             
-            self.device = "cpu"  # Force CPU usage for stability
+            self.device = "cpu"
             self._move_models_to_device()
             print("Models loaded successfully!")
         except Exception as e:
@@ -39,70 +38,54 @@ class FinancialAnalyzer:
         self.finbert_model.to(self.device)
         self.t5_model.to(self.device)
 
-    def process_file(self, file, file_type):
+    def read_file_content(self, file):
+        """Read and process uploaded file content"""
+        if file is None:
+            return "No file uploaded"
+            
         try:
-            if file_type == "csv":
+            if file.name.endswith('.csv'):
                 df = pd.read_csv(file.name)
                 return df.to_string()
-            elif file_type == "excel":
+            elif file.name.endswith(('.xls', '.xlsx')):
                 df = pd.read_excel(file.name)
                 return df.to_string()
-            elif file_type == "markdown":
+            elif file.name.endswith('.md'):
                 with open(file.name, 'r') as f:
                     return f.read()
+            else:
+                return "Unsupported file format. Please upload CSV, Excel, or Markdown files."
         except Exception as e:
             return f"Error processing file: {str(e)}"
 
-    def analyze_financials(self, balance_sheet_file, income_statement_file, file_type="csv"):
+    def analyze_financial_data(self, balance_sheet_file, income_statement_file):
+        """Analyze uploaded financial statements"""
         try:
-            # Process uploaded files
-            balance_sheet_data = self.process_file(balance_sheet_file, file_type)
-            income_statement_data = self.process_file(income_statement_file, file_type)
+            # Read file contents
+            balance_sheet = self.read_file_content(balance_sheet_file)
+            income_statement = self.read_file_content(income_statement_file)
+            
+            if "Error" in balance_sheet or "Error" in income_statement:
+                return "Error processing files. Please check the file format and content."
 
-            # Format the prompt for TinyLlama
-            prompt = self.format_financial_prompt(balance_sheet_data, income_statement_data)
+            # Format prompt for analysis
+            prompt = f"""<human>Analyze these financial statements:
             
-            # Generate insights using TinyLlama
-            insights = self.generate_insights(prompt)
+            Balance Sheet:
+            {balance_sheet[:1000]}
             
-            # Generate sentiment analysis
-            sentiment = self.analyze_sentiment(balance_sheet_data, income_statement_data)
+            Income Statement:
+            {income_statement[:1000]}
             
-            # Generate recommendations
-            recommendations = self.generate_recommendations(insights, sentiment)
-
-            # Combine results
-            analysis_results = {
-                "Financial Insights": insights,
-                "Sentiment Analysis": sentiment,
-                "Recommendations": recommendations
-            }
-
-            return json.dumps(analysis_results, indent=2)
-
-        except Exception as e:
-            return f"Error during analysis: {str(e)}"
-
-    def format_financial_prompt(self, balance_sheet, income_statement):
-        return f"""<human>Please analyze these financial statements and provide key insights:
-
-Balance Sheet Summary:
-{balance_sheet[:1000]}
-
-Income Statement Summary:
-{income_statement[:1000]}
-
-Please provide:
-1. Key financial metrics analysis
-2. Growth trends
-3. Risk factors
-4. Areas of concern
-5. Positive indicators</human>
-
-<assistant>I'll analyze the financial statements and provide comprehensive insights:"""
+            Provide:
+            1. Key financial metrics
+            2. Growth trends
+            3. Risk analysis
+            4. Recommendations</human>
+            
+            <assistant>Here's my analysis:"""
 
-    def generate_insights(self, prompt):
-        try:
+            # Generate analysis using TinyLlama
             inputs = self.tiny_tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
             outputs = self.tiny_model.generate(
                 inputs["input_ids"],
@@ -112,9 +95,21 @@ Please provide:
                 do_sample=True,
                 pad_token_id=self.tiny_tokenizer.eos_token_id
             )
-            return self.tiny_tokenizer.decode(outputs[0], skip_special_tokens=True)
+            analysis = self.tiny_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            # Generate sentiment
+            sentiment = self.analyze_sentiment(balance_sheet, income_statement)
+
+            # Format results
+            results = {
+                "Analysis": analysis,
+                "Sentiment": sentiment
+            }
+
+            return json.dumps(results, indent=2)
+
         except Exception as e:
-            return f"Error generating insights: {str(e)}"
+            return f"Error during analysis: {str(e)}"
 
     def analyze_sentiment(self, balance_sheet, income_statement):
         try:
@@ -125,58 +120,36 @@ Please provide:
             labels = ['negative', 'neutral', 'positive']
             return {
                 'sentiment': labels[probs.argmax().item()],
-                'confidence': f"{probs.max().item():.2f}",
-                'detailed_scores': {
-                    label: f"{prob:.2f}" 
-                    for label, prob in zip(labels, probs[0].tolist())
-                }
+                'confidence': f"{probs.max().item():.2f}"
             }
         except Exception as e:
             return f"Error in sentiment analysis: {str(e)}"
 
-    def generate_recommendations(self, insights, sentiment):
-        try:
-            prompt = f"summarize financial recommendations based on: {insights[:500]} Financial sentiment: {sentiment}"
-            inputs = self.t5_tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
-            outputs = self.t5_model.generate(
-                inputs["input_ids"],
-                max_length=200,
-                num_beams=4,
-                temperature=0.7,
-                top_p=0.95
-            )
-            return self.t5_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        except Exception as e:
-            return f"Error generating recommendations: {str(e)}"
-
-def create_gradio_interface():
+def create_interface():
     analyzer = FinancialAnalyzer()
     
-    def analyze_files(balance_sheet, income_statement, file_type):
-        return analyzer.analyze_financials(balance_sheet, income_statement, file_type)
-
     iface = gr.Interface(
-        fn=analyze_files,
+        fn=analyzer.analyze_financial_data,
         inputs=[
-            gr.File(label="Upload Balance Sheet"),
-            gr.File(label="Upload Income Statement"),
-            gr.Radio(
-                choices=["csv", "excel", "markdown"],
-                label="File Type",
-                value="csv"
+            gr.File(
+                label="Upload Balance Sheet (CSV, Excel, or Markdown)",
+                type="file"
+            ),
+            gr.File(
+                label="Upload Income Statement (CSV, Excel, or Markdown)",
+                type="file"
             )
         ],
-        outputs=gr.Textbox(label="Analysis Results", lines=20),
+        outputs=gr.Textbox(
+            label="Analysis Results",
+            lines=20
+        ),
         title="Financial Statement Analyzer",
-        description="Upload your financial statements to get AI-powered insights and recommendations.",
-        examples=[
-            ["balance_sheet.csv", "income_statement.csv", "csv"],
-            ["balance_sheet.xlsx", "income_statement.xlsx", "excel"],
-            ["balance_sheet.md", "income_statement.md", "markdown"]
-        ]
+        description="Upload your financial statements (Balance Sheet and Income Statement) to get AI-powered insights and analysis."
     )
+    
     return iface
 
 if __name__ == "__main__":
-    iface = create_gradio_interface()
+    iface = create_interface()
     iface.launch()