Update app.py

app.py

@@ -1,164 +1,166 @@
 import gradio as gr
 import pandas as pd
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import pipeline
 import torch
-
-def load_models():
-    """Initialize models with proper error handling"""
-    models = {}
-    try:
-        # 1. Load FinBERT (always works, no approval needed)
-        models['finbert'] = pipeline(
+import PyPDF2
+import io
+
+class FinancialAnalyzer:
+    def __init__(self):
+        """Initialize models"""
+        # 1. Llama 2 for strategic analysis
+        self.strategic_analyzer = pipeline(
+            "text-generation",
+            model="meta-llama/Llama-2-7b-chat-hf",
+            device_map="auto"
+        )
+        
+        # 2. FinBERT for financial sentiment
+        self.financial_analyzer = pipeline(
             "text-classification",
             model="ProsusAI/finbert",
             return_all_scores=True
         )
-        print("✓ FinBERT loaded successfully")
         
-        # 2. Try loading Llama 2 (requires approval)
-        try:
-            models['llama2'] = pipeline(
-                "text-generation",
-                model="meta-llama/Llama-2-7b-chat-hf",
-                torch_dtype=torch.float16,
-                device_map="auto"
-            )
-            print("✓ Llama 2 loaded successfully")
-        except Exception as e:
-            print(f"⚠️ Llama 2 not available: {str(e)}")
-            # Fallback to Falcon
-            models['llama2'] = pipeline(
-                "text-generation",
-                model="tiiuae/falcon-7b",
-                torch_dtype=torch.float16,
-                device_map="auto"
-            )
-            print("✓ Using Falcon as fallback for analysis")
-        
-        # 3. Load Falcon (always works, no approval needed)
-        models['falcon'] = pipeline(
+        # 3. Falcon for recommendations
+        self.recommendation_generator = pipeline(
             "text-generation",
-            model="tiiuae/falcon-7b",
-            torch_dtype=torch.float16,
+            model="tiiuae/falcon-7b-instruct",
             device_map="auto"
         )
-        print("✓ Falcon loaded successfully")
+    
+    def extract_text_from_pdf(self, pdf_file):
+        """Extract text from PDF file"""
+        try:
+            pdf_reader = PyPDF2.PdfReader(pdf_file)
+            text = ""
+            for page in pdf_reader.pages:
+                text += page.extract_text() + "\n"
+            return text
+        except Exception as e:
+            print(f"Error extracting PDF text: {str(e)}")
+            return None
+
+    def generate_strategic_analysis(self, financial_data, kpi_data=None):
+        """Generate strategic analysis using Llama 2"""
+        # Include KPI data if available
+        kpi_section = f"\nKPI Information:\n{kpi_data}" if kpi_data else ""
         
-        return models
-    except Exception as e:
-        print(f"Error loading models: {str(e)}")
-        return None
+        prompt = f"""[INST] As a senior financial analyst, analyze these financial statements and KPIs:
+
+        Financial Data:
+        {financial_data}
+        {kpi_section}
+
+        Provide a comprehensive analysis including:
+        1. Business Health Assessment
+        2. Key Strategic Insights
+        3. Market Position Analysis
+        4. Growth Opportunities
+        5. Risk Factors
+        6. KPI Performance and Trends
+        7. Financial Ratio Analysis
+        8. Operational Efficiency Assessment [/INST]"""
+        
+        response = self.strategic_analyzer(
+            prompt,
+            max_length=1500,  # Increased for more comprehensive analysis
+            temperature=0.7
+        )
+        return response[0]['generated_text']
+    
+    def analyze_sentiment(self, text):
+        """Analyze financial sentiment using FinBERT"""
+        return self.financial_analyzer(text)
+    
+    def generate_recommendations(self, analysis, kpi_data=None):
+        """Generate recommendations using Falcon"""
+        kpi_context = f"\nKPI Context:\n{kpi_data}" if kpi_data else ""
+        
+        prompt = f"""Based on this financial analysis and KPI data:
+        {analysis}
+        {kpi_context}
+        
+        Provide specific, actionable recommendations covering:
+        1. Strategic Initiatives
+        2. Operational Improvements
+        3. Financial Management
+        4. Risk Mitigation
+        5. Growth Strategy
+        6. KPI Optimization
+        7. Performance Enhancement
+        8. Resource Allocation"""
+        
+        response = self.recommendation_generator(
+            prompt,
+            max_length=1000,
+            temperature=0.6
+        )
+        return response[0]['generated_text']
 
-def analyze_financial_statements(income_statement, balance_sheet):
-    """Analyze financial statements using AI models"""
+def analyze_financial_statements(income_statement, balance_sheet, kpi_pdf=None):
+    """Main analysis function"""
     try:
-        # Load data
+        # Read financial statements
         income_df = pd.read_csv(income_statement.name)
         balance_df = pd.read_csv(balance_sheet.name)
         
-        # Prepare financial data summary
-        financial_data = prepare_financial_summary(income_df, balance_df)
+        # Initialize analyzer
+        analyzer = FinancialAnalyzer()
         
-        # Load models
-        models = load_models()
-        if not models:
-            return "Error: Could not load analysis models."
+        # Extract KPI data if provided
+        kpi_data = None
+        if kpi_pdf:
+            kpi_data = analyzer.extract_text_from_pdf(kpi_pdf.name)
         
-        # Generate analyses
-        analysis = generate_analysis(models, financial_data)
+        # Prepare financial data
+        financial_data = f"""
+        Income Statement Summary:
+        {income_df.to_string()}
         
-        # Format and return results
-        return format_results(analysis)
+        Balance Sheet Summary:
+        {balance_df.to_string()}
+        """
+        
+        # Generate strategic analysis
+        strategic_analysis = analyzer.generate_strategic_analysis(financial_data, kpi_data)
+        
+        # Analyze sentiment
+        sentiment = analyzer.analyze_sentiment(strategic_analysis)
+        
+        # Generate recommendations
+        recommendations = analyzer.generate_recommendations(strategic_analysis, kpi_data)
+        
+        # Format output
+        output = format_results(strategic_analysis, sentiment, recommendations, kpi_data)
+        
+        return output
         
     except Exception as e:
         return f"Error analyzing files: {str(e)}"
 
-def prepare_financial_summary(income_df, balance_df):
-    """Prepare financial data summary"""
-    latest_year = income_df['Period'].iloc[-1]
-    previous_year = income_df['Period'].iloc[-2]
-    
-    summary = f"""Financial Analysis for {latest_year}:
-
-Income Statement Summary:
-- Revenue: {income_df['Revenue'].iloc[-1]}
-- Gross Profit: {income_df['Gross Profit'].iloc[-1]}
-- Net Income: {income_df['Net Income'].iloc[-1]}
-- YoY Revenue Growth: {(income_df['Revenue'].iloc[-1] / income_df['Revenue'].iloc[-2] - 1) * 100:.1f}%
-- YoY Net Income Growth: {(income_df['Net Income'].iloc[-1] / income_df['Net Income'].iloc[-2] - 1) * 100:.1f}%
-
-Balance Sheet Summary:
-- Total Assets: {balance_df['Total Assets'].iloc[-1]}
-- Total Liabilities: {balance_df['Total Liabilities'].iloc[-1]}
-- Shareholder's Equity: {balance_df["Shareholder's Equity"].iloc[-1]}
-- Current Ratio: {balance_df['Total current assets'].iloc[-1] / balance_df['Total current liabilities'].iloc[-1]:.2f}"""
-    
-    return summary
-
-def generate_analysis(models, financial_data):
-    """Generate comprehensive analysis using all models"""
-    
-    # 1. Strategic Analysis (Llama 2 or Falcon)
-    strategy_prompt = f"""[INST] As a senior financial analyst, analyze these financial metrics and provide strategic insights:
-
-{financial_data}
-
-Provide analysis covering:
-1. Financial Health Assessment
-2. Key Performance Indicators
-3. Growth Analysis
-4. Risk Factors
-5. Strategic Position [/INST]"""
-
-    strategic_analysis = models['llama2'](
-        strategy_prompt,
-        max_length=1000,
-        temperature=0.7
-    )[0]['generated_text']
-
-    # 2. Financial Sentiment (FinBERT)
-    sentiment = models['finbert'](strategic_analysis)
-
-    # 3. Recommendations (Falcon)
-    recommendations_prompt = f"""Based on this financial analysis:
-{strategic_analysis}
-
-Provide specific, actionable recommendations for:
-1. Immediate Actions (0-6 months)
-2. Short-term Strategy (6-12 months)
-3. Medium-term Initiatives (1-2 years)
-4. Risk Mitigation
-5. Growth Opportunities"""
-
-    recommendations = models['falcon'](
-        recommendations_prompt,
-        max_length=800,
-        temperature=0.6
-    )[0]['generated_text']
-
-    return {
-        'strategic_analysis': strategic_analysis,
-        'sentiment': sentiment,
-        'recommendations': recommendations
-    }
-
-def format_results(analysis):
-    """Format analysis results into readable report"""
+def format_results(analysis, sentiment, recommendations, kpi_data=None):
+    """Format analysis results"""
     output = "# Financial Analysis Report\n\n"
     
     # Strategic Analysis
     output += "## Strategic Analysis\n\n"
-    output += analysis['strategic_analysis'].split('[/INST]')[-1].strip() + "\n\n"
+    output += analysis.split('[/INST]')[-1].strip() + "\n\n"
+    
+    # KPI Analysis (if available)
+    if kpi_data:
+        output += "## KPI Analysis\n\n"
+        output += "Analysis includes KPI data from provided documentation.\n\n"
     
-    # Market Sentiment
+    # Sentiment Analysis
     output += "## Market Sentiment\n\n"
-    for score in analysis['sentiment'][0]:
+    for score in sentiment[0]:
         output += f"- {score['label']}: {score['score']:.2%}\n"
     output += "\n"
     
     # Recommendations
     output += "## Strategic Recommendations\n\n"
-    output += analysis['recommendations']
+    output += recommendations
     
     return output
 
@@ -167,20 +169,25 @@ iface = gr.Interface(
     fn=analyze_financial_statements,
     inputs=[
         gr.File(label="Income Statement (CSV)"),
-        gr.File(label="Balance Sheet (CSV)")
+        gr.File(label="Balance Sheet (CSV)"),
+        gr.File(label="KPI Documentation (PDF, Optional)", file_types=[".pdf"])
     ],
     outputs=gr.Markdown(),
     title="AI-Powered Financial Statement Analysis",
-    description="""Upload your financial statements for comprehensive analysis using:
-    - Llama 2 / Falcon: Strategic Analysis
+    description="""Upload your financial statements and optional KPI documentation for comprehensive analysis using:
+    - Llama 2: Strategic Analysis
     - FinBERT: Financial Sentiment Analysis
     - Falcon: Strategic Recommendations
     
-    Note: The system will automatically use available models based on access permissions.""",
+    Supports:
+    - Income Statement (CSV)
+    - Balance Sheet (CSV)
+    - KPI Documentation (PDF)""",
     examples=[
         [
             "OFINTECH-Income Statement-template.csv",
-            "OFINTECH Balance Sheet template.csv"
+            "OFINTECH Balance Sheet template.csv",
+            None  # Optional KPI PDF
         ]
     ]
 )