Update app.py

app.py

@@ -1,5 +1,3 @@
-
-
 import torch
 import yfinance as yf
 import matplotlib.pyplot as plt
@@ -15,13 +13,16 @@ import spaces
 import pandas as pd
 
 
+
+
 # Configure logging
 logging.basicConfig(filename='debug.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Load the model and processor
+# Load the chart_analysis model and processor
 processor = AutoProcessor.from_pretrained("mobenta/chart_analysis")
 model = AutoModelForPreTraining.from_pretrained("mobenta/chart_analysis")
 
+@spaces.GPU
 def predict(image, input_text):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
@@ -29,7 +30,7 @@ def predict(image, input_text):
     image = image.convert("RGB")
     inputs = processor(text=input_text, images=image, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
-
+    
     prompt_length = inputs['input_ids'].shape[1]
     generate_ids = model.generate(**inputs, max_new_tokens=512)
     output_text = processor.batch_decode(generate_ids[:, prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
@@ -136,23 +137,47 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
         resized_image = image.resize(new_size, Image.LANCZOS)
         resized_image.save(filename)
 
-        logging.debug(f"Resized image saved to {filename}")
-        return filename
+        logging.debug(f"Resized image with timeframe {timeframe} and ticker {ticker} saved to {filename}")
     except Exception as e:
-        logging.error(f"Error creating stock chart: {e}")
+        logging.error(f"Error creating or resizing chart: {e}")
+        raise
+
+def combine_images(image_paths, output_path='combined_chart.png'):
+    try:
+        logging.debug(f"Combining images {image_paths} into {output_path}")
+        images = [Image.open(path) for path in image_paths]
+
+        # Calculate total width and max height for combined image
+        total_width = sum(img.width for img in images)
+        max_height = max(img.height for img in images)
+
+        combined_image = Image.new('RGB', (total_width, max_height))
+        x_offset = 0
+        for img in images:
+            combined_image.paste(img, (x_offset, 0))
+            x_offset += img.width
+
+        combined_image.save(output_path)
+        logging.debug(f"Combined image saved to {output_path}")
+        return output_path
+    except Exception as e:
+        logging.error(f"Error combining images: {e}")
         raise
 
 def gradio_interface(ticker1, ticker2, ticker3, ticker4, start_date, end_date, query, analysis_type, interval, indicators):
     try:
-        chart_paths = []
+        logging.debug(f"Starting gradio_interface with tickers: {ticker1}, {ticker2}, {ticker3}, {ticker4}, start_date: {start_date}, end_date: {end_date}, query: {query}, analysis_type: {analysis_type}, interval: {interval}")
+
         tickers = [ticker1, ticker2, ticker3, ticker4]
+        chart_paths = []
 
-        for ticker in tickers:
-            if ticker.strip():
-                data = fetch_stock_data(ticker, start_date, end_date, interval)
-                chart_path = f"{ticker}_chart.png"
-                create_stock_chart(data, ticker, filename=chart_path, timeframe=interval, indicators=indicators)
-                chart_paths.append(chart_path)
+        for i, ticker in enumerate(tickers):
+            if ticker:
+                data = fetch_stock_data(ticker, start=start_date, end=end_date, interval=interval)
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_chart:
+                    chart_path = temp_chart.name
+                    create_stock_chart(data, ticker, chart_path, timeframe=interval, indicators=indicators)
+                    chart_paths.append(chart_path)
 
         if analysis_type == 'Comparative Analysis' and len(chart_paths) > 1:
             with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_combined_chart:
@@ -161,6 +186,7 @@ def gradio_interface(ticker1, ticker2, ticker3, ticker4, start_date, end_date, q
                 insights = predict(Image.open(combined_chart_path), query)
                 return insights, combined_chart_path
 
+        # No comparative analysis, just return the single chart
         if chart_paths:
             insights = predict(Image.open(chart_paths[0]), query)
             return insights, chart_paths[0]
@@ -172,23 +198,7 @@ def gradio_interface(ticker1, ticker2, ticker3, ticker4, start_date, end_date, q
 
 def gradio_app():
     with gr.Blocks() as demo:
-        gr.Markdown("""
-        ## 📈Stock Analysis Dashboard
-
-        This application provides a comprehensive stock analysis tool that allows users to input up to four stock tickers, specify date ranges, and select various financial indicators. The core functionalities include:
-
-        1. **Data Fetching and Chart Creation**: Historical stock data is fetched from Yahoo Finance, and candlestick charts are generated with optional financial indicators like RSI, SMA, VWAP, and Bollinger Bands.
-
-        2. **Text Analysis and Insights Generation**: The application uses a pre-trained model based on the **Paligema** architecture to analyze the input chart and text query, generating insightful analysis based on the provided financial data and context.
-
-        3. **User Interface**: Users can interactively select stocks, date ranges, intervals, and indicators. The app also supports the analysis of single tickers or comparative analysis across multiple tickers.
-
-        4. **Logging and Debugging**: Detailed logging helps in debugging and tracking the application's processes.
-
-        5. **Enhanced Image Processing**: The app adds financial metrics and annotations to the generated charts, ensuring clear presentation of data.
-
-        This tool leverages the Paligema model to provide detailed insights into stock market trends, offering an interactive and educational experience for users.
-        """)
+        gr.Markdown("## Stock Analysis Dashboard")
 
         with gr.Row():
             ticker1 = gr.Textbox(label="Primary Ticker", value="GC=F")
@@ -220,4 +230,3 @@ def gradio_app():
 
 if __name__ == "__main__":
     gradio_app()
-