Update app.py

app.py

@@ -3,10 +3,8 @@ import gradio as gr
 import pandas as pd
 import torch
 import logging
-import gc
-import threading
-import concurrent.futures
 from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
+import gc
 
 # Setup logging
 logging.basicConfig(
@@ -25,452 +23,193 @@ def clear_gpu_memory():
         torch.cuda.empty_cache()
     gc.collect()
 
-class ModelLoadingError(Exception):
-    """Custom exception for model loading errors"""
-    pass
-
-class ModelManager:
-    """Handles model loading and inference"""
-    
-    def __init__(self):
-        self.device = DEVICE
-        self.models = {}
-        self.tokenizers = {}
-        self.model_cache = {}
-        self.max_cache_size = 2
-    
-    def load_model(self, model_name, model_type="sentiment", timeout=300):
-        """Load model and tokenizer with thread-safe timeout"""
-        try:
-            if model_name in self.model_cache:
-                self.models[model_name] = self.model_cache[model_name]
-                logger.info(f"Loaded {model_name} from cache")
-                return
-
-            def load_model_task():
-                if model_type == "sentiment":
-                    self.tokenizers[model_name] = AutoTokenizer.from_pretrained(
-                        model_name,
-                        use_fast=True
-                    )
-                    self.models[model_name] = AutoModelForSequenceClassification.from_pretrained(
-                        model_name,
-                        torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
-                    ).to(self.device)
-                else:
-                    self.models[model_name] = pipeline(
-                        "text-generation",
-                        model=model_name,
-                        device_map="auto" if self.device == "cuda" else None,
-                        torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
-                    )
-
-            # Use ThreadPoolExecutor for timeout
-            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
-                future = executor.submit(load_model_task)
-                try:
-                    future.result(timeout=timeout)
-                except concurrent.futures.TimeoutError:
-                    raise ModelLoadingError(f"Model loading timed out after {timeout} seconds")
-
-            # Cache the model
-            self.cache_model(model_name, self.models[model_name])
-            logger.info(f"Successfully loaded model: {model_name}")
-
-        except Exception as e:
-            logger.error(f"Error loading model {model_name}: {str(e)}")
-            raise ModelLoadingError(f"Failed to load model {model_name}: {str(e)}")
-
-    def cache_model(self, model_name, model):
-        """Cache model for faster reloading"""
-        if len(self.model_cache) >= self.max_cache_size:
-            oldest_model = next(iter(self.model_cache))
-            del self.model_cache[oldest_model]
-        self.model_cache[model_name] = model
-
-    def unload_model(self, model_name):
-        """Unload model and tokenizer"""
-        try:
-            if model_name in self.models:
-                del self.models[model_name]
-            if model_name in self.tokenizers:
-                del self.tokenizers[model_name]
-            clear_gpu_memory()
-            logger.info(f"Unloaded model: {model_name}")
-        except Exception as e:
-            logger.error(f"Error unloading model {model_name}: {str(e)}")
-
-    def get_model(self, model_name):
-        """Get loaded model"""
-        return self.models.get(model_name)
-
-    def get_tokenizer(self, model_name):
-        """Get loaded tokenizer"""
-        return self.tokenizers.get(model_name)
-
 class FinancialAnalyzer:
-    """Main analyzer class for financial statements"""
+    """Simplified Financial Analyzer using small models"""
     
     def __init__(self):
-        self.model_manager = ModelManager()
-        self.models = {
-            "sentiment": "ProsusAI/finbert",
-            "analysis": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-            "recommendation": "tiiuae/falcon-rw-1b"
-        }
-        
-        # Load sentiment model at initialization with longer timeout
+        # Initialize with two small models
+        self.sentiment_model = None
+        self.analysis_model = None
+        self.load_models()
+
+    def load_models(self):
+        """Load the required models"""
         try:
-            self.model_manager.load_model(self.models["sentiment"], "sentiment", timeout=600)
-            logger.info("Sentiment model initialized successfully")
+            # Load FinBERT for sentiment analysis
+            self.sentiment_model = pipeline(
+                "text-classification",
+                model="ProsusAI/finbert",
+                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
+            )
+            
+            # Load small model for analysis and recommendations
+            self.analysis_model = pipeline(
+                "text-generation",
+                model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
+            )
+            
+            logger.info("Models loaded successfully")
         except Exception as e:
-            logger.error(f"Failed to initialize sentiment model: {str(e)}")
+            logger.error(f"Error loading models: {str(e)}")
             raise
 
-    def read_csv(self, file_obj):
-        """Read and validate CSV file with better error handling"""
+    def process_csv(self, file_obj):
+        """Process CSV file and extract KPIs"""
         try:
             if file_obj is None:
                 raise ValueError("No file provided")
             
-            # Read CSV with explicit encoding and error handling
-            df = pd.read_csv(file_obj, encoding='utf-8', on_bad_lines='skip')
+            df = pd.read_csv(file_obj)
             
             if df.empty:
                 raise ValueError("Empty CSV file")
-            
-            # Log CSV information
-            logger.info(f"CSV Preview:\n{df.head()}")
-            logger.info(f"CSV Columns: {df.columns.tolist()}")
-            
-            # Validate numeric columns
+
+            # Get numeric columns
             numeric_cols = df.select_dtypes(include=['float64', 'int64']).columns
             if len(numeric_cols) == 0:
                 raise ValueError("No numeric columns found in CSV")
-            
-            # Generate statistical summary
+
+            # Calculate basic KPIs
             summary = df[numeric_cols].describe()
-            logger.info(f"Statistical Summary:\n{summary}")
             
-            return summary
+            # Extract key metrics
+            metrics = {
+                'total': df[numeric_cols].sum(),
+                'average': df[numeric_cols].mean(),
+                'growth': df[numeric_cols].pct_change().mean() * 100
+            }
+            
+            return summary, metrics
             
         except Exception as e:
-            logger.error(f"Error reading CSV: {str(e)}")
+            logger.error(f"Error processing CSV: {str(e)}")
             raise
 
-    def analyze_sentiment(self, text):
-        """Analyze sentiment using FinBERT with improved error handling"""
-        try:
-            model_name = self.models["sentiment"]
-            model = self.model_manager.get_model(model_name)
-            tokenizer = self.model_manager.get_tokenizer(model_name)
-
-            # Validate input
-            if not text or not isinstance(text, str):
-                raise ValueError("Invalid input text")
-
-            # Preprocess text
-            text = text.strip()
-            if len(text) == 0:
-                raise ValueError("Empty text input")
-
-            # Truncate text if too long
-            max_length = 512
-            if len(text.split()) > max_length:
-                logger.warning(f"Text length exceeds {max_length} tokens. Truncating...")
-
-            # Tokenize with proper padding and truncation
-            inputs = tokenizer(
-                text,
-                return_tensors="pt",
-                truncation=True,
-                max_length=max_length,
-                padding=True
-            ).to(DEVICE)
-
-            # Get prediction
-            with torch.no_grad():
-                outputs = model(**inputs)
-                probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
-
-            # Process results
-            labels = ['negative', 'neutral', 'positive']
-            scores = probabilities[0].cpu().tolist()
-
-            results = [
-                {'label': label, 'score': score}
-                for label, score in zip(labels, scores)
-            ]
-
-            logger.info(f"Sentiment analysis results: {results}")
-            return [results]
-
-        except Exception as e:
-            logger.error(f"Sentiment analysis error: {str(e)}")
-            return [{"label": "error", "score": 1.0}]
-
-    def generate_analysis(self, financial_data):
-        """Generate strategic analysis with improved prompting and error handling"""
-        try:
-            model_name = self.models["analysis"]
-            self.model_manager.load_model(model_name, "generation", timeout=600)
-
-            # Truncate financial data if too long
-            max_data_length = 1000
-            if len(financial_data.split()) > max_data_length:
-                logger.warning(f"Financial data too long. Truncating to {max_data_length} tokens...")
-                financial_data = ' '.join(financial_data.split()[:max_data_length])
-
-            prompt = f"""[INST] As a senior financial analyst, provide a detailed analysis of these financial statements:
-
-            Financial Data:
-            {financial_data}
-
-            Please provide a comprehensive analysis covering:
-
-            1. Business Health Assessment
-            - Current financial position
-            - Key performance indicators
-            - Trend analysis
-
-            2. Key Strategic Insights
-            - Major financial trends
-            - Performance drivers
-            - Areas of concern
-
-            3. Market Position
-            - Competitive advantages
-            - Market share indicators
-            - Industry comparison
-
-            4. Growth Opportunities
-            - Expansion potential
-            - Investment opportunities
-            - Revenue growth areas
-
-            5. Risk Factors
-            - Financial risks
-            - Operational risks
-            - Market risks
-
-            Provide specific metrics and detailed explanations for each section. [/INST]"""
-
-            logger.info("Generating analysis...")
-            response = self.model_manager.get_model(model_name)(
-                prompt,
-                max_length=2000,
-                min_length=800,
-                temperature=0.7,
-                do_sample=True,
-                num_return_sequences=1,
-                truncation=True,
-                repetition_penalty=1.2,
-                no_repeat_ngram_size=3
-            )
-
-            analysis_text = response[0]['generated_text']
-            return self.format_analysis_text(analysis_text)
-
-        except Exception as e:
-            logger.error(f"Analysis generation error: {str(e)}")
-            return "Error in analysis generation"
-        finally:
-            self.model_manager.unload_model(model_name)
-
-    def format_analysis_text(self, text):
-        """Format the analysis text for better readability"""
+    def analyze_financials(self, income_summary, balance_summary):
+        """Generate financial analysis and recommendations"""
         try:
-            sections = text.split('\n\n')
-            formatted_sections = []
+            financial_context = f"""
+            Income Statement Metrics:
+            {income_summary[0].to_string()}
             
-            current_section = None
-            for section in sections:
-                section = section.strip()
-                if not section:
-                    continue
-                
-                # Check if this is a new section
-                if any(section.startswith(str(i)) for i in range(1, 6)):
-                    current_section = f"### {section}"
-                    formatted_sections.append(current_section)
-                elif current_section:
-                    # Add bullet points to content under sections
-                    lines = section.split('\n')
-                    formatted_lines = []
-                    for line in lines:
-                        line = line.strip()
-                        if line:
-                            if not line.startswith('- '):
-                                line = f"- {line}"
-                            formatted_lines.append(line)
-                    formatted_sections.append('\n'.join(formatted_lines))
+            Key Income Indicators:
+            {income_summary[1]}
             
-            return '\n\n'.join(formatted_sections)
-        except Exception as e:
-            logger.error(f"Error formatting analysis text: {str(e)}")
-            return text
-
-    def generate_recommendations(self, analysis):
-        """Generate recommendations with improved prompting and error handling"""
-        try:
-            model_name = self.models["recommendation"]
-            self.model_manager.load_model(model_name, "generation", timeout=600)
-
-            # Truncate analysis if too long
-            max_analysis_length = 1000
-            if len(analysis.split()) > max_analysis_length:
-                logger.warning(f"Analysis too long. Truncating to {max_analysis_length} tokens...")
-                analysis = ' '.join(analysis.split()[:max_analysis_length])
-
-            prompt = f"""Based on this financial analysis, provide detailed strategic recommendations:
-
-            Analysis Context:
-            {analysis}
-
-            Please provide specific, actionable recommendations for each area:
-
-            1. Strategic Initiatives
-            - Detail specific actions for business growth
-            - Identify market expansion opportunities
-            - Outline product/service development strategies
-
-            2. Operational Improvements
-            - Specify efficiency enhancement measures
-            - Recommend process optimization steps
-            - Suggest cost reduction strategies
-
-            3. Financial Management
-            - Provide cash flow optimization tactics
-            - Prioritize investment opportunities
-            - Detail risk management approaches
+            Balance Sheet Metrics:
+            {balance_summary[0].to_string()}
+            
+            Key Balance Sheet Indicators:
+            {balance_summary[1]}
+            """
 
-            4. Risk Mitigation
-            - Address identified risks
-            - Outline specific mitigation strategies
-            - Suggest monitoring mechanisms
+            # Generate sentiment analysis
+            sentiment = self.sentiment_model(financial_context)[0]
+            
+            # Generate business analysis
+            analysis_prompt = f"""[INST] Based on the following financial data, provide:
+            1. Current Business Status
+            2. Key Business Insights
+            3. Strategic Recommendations and Roadmap
 
-            5. Growth Strategy
-            - Identify market opportunities
-            - Detail expansion plans
-            - Specify resource requirements
+            Financial Context:
+            {financial_context}
 
-            For each recommendation:
-            - Include implementation timeline
-            - Specify resource requirements
-            - Define success metrics
-            - List potential challenges
+            Sentiment: {sentiment['label']} ({sentiment['score']:.2%})
 
-            Format each section with clear, actionable bullet points."""
+            Provide a concise but detailed analysis for each section.
+            [/INST]"""
 
-            logger.info("Generating recommendations...")
-            response = self.model_manager.get_model(model_name)(
-                prompt,
-                max_length=2000,
-                min_length=800,
+            response = self.analysis_model(
+                analysis_prompt,
+                max_length=1000,
                 temperature=0.7,
-                do_sample=True,
-                num_return_sequences=1,
-                truncation=True,
-                repetition_penalty=1.2,
-                no_repeat_ngram_size=3
+                num_return_sequences=1
             )
 
-            recommendations_text = response[0]['generated_text']
-            return self.format_recommendation_text(recommendations_text)
+            return self.format_response(response[0]['generated_text'], sentiment)
 
         except Exception as e:
-            logger.error(f"Recommendations generation error: {str(e)}")
-            return "Error generating recommendations"
-        finally:
-            self.model_manager.unload_model(model_name)
+            logger.error(f"Error in analysis: {str(e)}")
+            return "Error generating analysis"
 
-    def format_recommendation_text(self, text):
-        """Format the recommendation text for better readability"""
+    def format_response(self, analysis_text, sentiment):
+        """Format the analysis response into structured sections"""
         try:
-            sections = text.split('\n\n')
-            formatted_sections = []
+            # Split the analysis into sections
+            sections = analysis_text.split('\n\n')
+            
+            # Initialize output sections
+            status = []
+            insights = []
+            recommendations = []
             
+            # Process each section
             current_section = None
             for section in sections:
-                section = section.strip()
-                if not section:
-                    continue
-                
-                # Check if this is a new section
-                if any(section.startswith(str(i)) for i in range(1, 6)):
-                    current_section = f"### {section}"
-                    formatted_sections.append(current_section)
-                elif current_section:
-                    # Add bullet points to content under sections
-                    lines = section.split('\n')
-                    formatted_lines = []
-                    for line in lines:
-                        line = line.strip()
-                        if line:
-                            if not line.startswith('- '):
-                                line = f"- {line}"
-                            formatted_lines.append(line)
-                    formatted_sections.append('\n'.join(formatted_lines))
+                if "Business Status" in section:
+                    current_section = status
+                elif "Key Business Insights" in section:
+                    current_section = insights
+                elif "Strategic Recommendations" in section:
+                    current_section = recommendations
+                elif current_section is not None:
+                    current_section.append(section.strip())
+            
+            # Format the final output
+            output = [
+                "# Financial Analysis Report\n\n",
+                f"## Overall Sentiment: {sentiment['label'].upper()} ({sentiment['score']:.2%})\n\n",
+                "## Current Business Status\n",
+                "".join(f"- {item}\n" for item in status if item),
+                "\n## Key Business Insights\n",
+                "".join(f"- {item}\n" for item in insights if item),
+                "\n## Strategic Recommendations & Roadmap\n",
+                "".join(f"- {item}\n" for item in recommendations if item)
+            ]
+            
+            return "".join(output)
             
-            return '\n\n'.join(formatted_sections)
         except Exception as e:
-            logger.error(f"Error formatting recommendation text: {str(e)}")
-            return text
+            logger.error(f"Error formatting response: {str(e)}")
+            return "Error formatting analysis results"
 
-def analyze_financial_statements(income_statement, balance_sheet):
-    """Main analysis function with improved error handling and logging"""
+def analyze_statements(income_statement, balance_sheet):
+    """Main function to analyze financial statements"""
     try:
-        clear_gpu_memory()
-        logger.info("Starting financial analysis...")
-        analyzer = FinancialAnalyzer()
-
-        # Validate inputs
-        if not income_statement or not balance_sheet:
-            return "Error: Please provide both income statement and balance sheet files"
+        # Check if files are uploaded
+        if income_statement is None or balance_sheet is None:
+            return "Please upload both Income Statement and Balance Sheet CSV files."
 
-        # Process financial statements
-        logger.info("Processing income statement...")
-        income_summary = analyzer.read_csv(income_statement)
-        logger.info("Processing balance sheet...")
-        balance_summary = analyzer.read_csv(balance_sheet)
+        # Get file names
+        income_filename = income_statement.name if hasattr(income_statement, 'name') else 'Income Statement'
+        balance_filename = balance_sheet.name if hasattr(balance_sheet, 'name') else 'Balance Sheet'
+        
+        logger.info(f"Processing {income_filename} and {balance_filename}")
 
-        financial_data = f"""
-        Income Statement Summary:
-        {income_summary.to_string()}
+        # Initialize analyzer
+        analyzer = FinancialAnalyzer()
         
-        Balance Sheet Summary:
-        {balance_summary.to_string()}
-        """
+        # Process statements with better error handling
+        try:
+            income_summary = analyzer.process_csv(income_statement)
+            logger.info("Successfully processed Income Statement")
+        except Exception as e:
+            return f"Error processing Income Statement: {str(e)}\nPlease ensure it's a valid CSV file with numeric data."
 
+        try:
+            balance_summary = analyzer.process_csv(balance_sheet)
+            logger.info("Successfully processed Balance Sheet")
+        except Exception as e:
+            return f"Error processing Balance Sheet: {str(e)}\nPlease ensure it's a valid CSV file with numeric data."
+        
         # Generate analysis
-        logger.info("Starting strategic analysis generation...")
-        analysis = analyzer.generate_analysis(financial_data)
-        if "Error" in analysis:
-            logger.error("Strategic analysis generation failed")
-            return "Error: Failed to generate strategic analysis. Please try again."
-
-        # Analyze sentiment
-        logger.info("Starting sentiment analysis...")
-        sentiment = analyzer.analyze_sentiment(analysis)
-        if sentiment[0][0]['label'] == "error":
-            logger.error("Sentiment analysis failed")
-            return "Error: Failed to analyze sentiment. Please try again."
-
-        # Generate recommendations
-        logger.info("Starting recommendations generation...")
-        recommendations = analyzer.generate_recommendations(analysis)
-        if "Error" in recommendations:
-            logger.error("Recommendations generation failed")
-            return "Error: Failed to generate recommendations. Please try again."
-
-        # Format results
-        logger.info("Formatting final results...")
-        result = format_results(analysis, sentiment, recommendations)
-        clear_gpu_memory()
+        logger.info("Generating analysis...")
+        result = analyzer.analyze_financials(income_summary, balance_summary)
         
-        logger.info("Analysis completed successfully")
+        clear_gpu_memory()
         return result
-
+        
     except Exception as e:
         logger.error(f"Analysis error: {str(e)}")
         return f"""Analysis Error:
@@ -478,99 +217,58 @@ def analyze_financial_statements(income_statement, balance_sheet):
         {str(e)}
         
         Please verify:
-        1. Files are valid CSV format
-        2. Files contain required financial data
-        3. File size is within limits (max 10MB)
-        4. Data contains numeric columns
-        5. Files are not corrupted"""
-
-def format_results(analysis, sentiment, recommendations):
-    """Format analysis results with improved validation and formatting"""
-    try:
-        if not isinstance(analysis, str) or not isinstance(recommendations, str):
-            raise ValueError("Invalid input types")
-
-        output = [
-            "# Financial Analysis Report\n\n",
-            "## Strategic Analysis\n\n",
-            f"{analysis.strip()}\n\n",
-            "## Market Sentiment\n\n"
-        ]
-
-        if isinstance(sentiment, list) and sentiment:
-            for score in sentiment[0]:
-                if isinstance(score, dict) and 'label' in score and 'score' in score:
-                    output.append(f"- {score['label']}: {score['score']:.2%}\n")
-        output.append("\n")
-
-        output.append("## Strategic Recommendations\n\n")
-        output.append(f"{recommendations.strip()}")
+        1. Files are in CSV format
+        2. Files contain numeric data columns
+        3. Files follow standard financial statement format"""
 
-        return "".join(output)
-    except Exception as e:
-        logger.error(f"Formatting error: {str(e)}")
-        return "Error formatting results"
-
-# Create Gradio interface with improved error handling and guidance
+# Create Gradio interface with improved file handling
 iface = gr.Interface(
-    fn=analyze_financial_statements,
+    fn=analyze_statements,
     inputs=[
         gr.File(
-            label="Income Statement (CSV)",
-            file_types=[".csv"]
+            label="Upload Income Statement (CSV)",
+            file_types=[".csv"],
+            file_count="single"
         ),
         gr.File(
-            label="Balance Sheet (CSV)",
-            file_types=[".csv"]
+            label="Upload Balance Sheet (CSV)",
+            file_types=[".csv"],
+            file_count="single"
         )
     ],
     outputs=gr.Markdown(),
-    title="AI-Powered Financial Statement Analyzer",
-    description="""## Financial Statement Analysis Tool
+    title="Financial Statement Analyzer",
+    description="""## Financial Analysis Tool
+
+How to use:
+1. Click 'Upload Income Statement' to select your income statement CSV file
+2. Click 'Upload Balance Sheet' to select your balance sheet CSV file
+3. Wait for the analysis to complete
 
-This tool provides comprehensive financial analysis using advanced AI models:
-- Strategic Analysis: In-depth analysis of financial position and trends
-- Sentiment Analysis: Assessment of financial health sentiment
-- Strategic Recommendations: Actionable insights and recommendations
+The tool will provide:
+- Business Status Assessment
+- Key Financial Insights
+- Strategic Recommendations
 
 Requirements:
 - Files must be in CSV format
 - Must contain numeric data columns
-- Maximum file size: 10MB
-- Standard financial statement format preferred
-
-Note: Analysis may take a few minutes to complete.""",
-    article="""### Usage Tips:
-1. Ensure your CSV files have clear column headers
-2. Verify that numeric data is properly formatted
-3. Wait for the analysis to complete - it may take several minutes
-4. The more detailed your financial data, the better the analysis
-
-For optimal results, include key financial metrics such as:
-- Revenue
-- Expenses
-- Profits/Losses
-- Assets
-- Liabilities
-- Equity""",
+- Standard financial statement format preferred""",
     examples=[
         ["example_income_statement.csv", "example_balance_sheet.csv"]
     ],
     flagging_mode="never"
 )
 
-# Launch the interface with proper error handling
+# Launch the interface with better error handling
 if __name__ == "__main__":
     try:
-        # Enable queue for better handling of multiple requests
-        iface.queue()
-        
-        # Launch with specific server configuration
+        iface.queue()  # Enable queuing for better file handling
         iface.launch(
             share=False,
             server_name="0.0.0.0",
             server_port=7860,
-            show_error=True
+            show_api=False  # Disable API tab for security
         )
     except Exception as e:
         logger.error(f"Launch error: {str(e)}")