Create app.py

app.py

@@ -0,0 +1,241 @@
+import torch
+import yfinance as yf
+import matplotlib.pyplot as plt
+import mplfinance as mpf
+from PIL import Image
+import gradio as gr
+import logging
+
+
+import torch
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+from PIL import Image
+import gradio as gr
+
+# Load the ChartGemma model and processor
+model = PaliGemmaForConditionalGeneration.from_pretrained("ahmed-masry/chartgemma", torch_dtype=torch.float16)
+processor = AutoProcessor.from_pretrained("ahmed-masry/chartgemma")
+
+
+# Configure logging to write to a file
+logging.basicConfig(filename='debug.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Use GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# model = YourModel.from_pretrained('your-model') # Uncomment and update this line with the actual model loading
+# processor = YourProcessor.from_pretrained('your-processor') # Uncomment and update this line with the actual processor loading
+# model = model.to(device)
+
+# Function to fetch stock data with different intervals
+def fetch_stock_data(ticker='TSLA', start='2023-01-01', end='2024-01-01', interval='1d'):
+    try:
+        logging.debug(f"Fetching data for {ticker} from {start} to {end} with interval {interval}")
+        stock = yf.Ticker(ticker)
+        data = stock.history(start=start, end=end, interval=interval)
+        logging.debug(f"Fetched data with {len(data)} rows")
+        return data
+    except Exception as e:
+        logging.error(f"Error fetching data: {e}")
+        raise
+
+# Function to create a candlestick chart with increased size and add timeframe and ticker information
+def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d'):
+    try:
+        logging.debug(f"Creating chart for {ticker} with timeframe {timeframe} and saving to {filename}")
+        title = f"{ticker.upper()} Price Data (Timeframe: {timeframe})"
+        
+        # Set the title font size using rcParams
+        plt.rcParams["axes.titlesize"] = 10
+        
+        # Define the style for the chart
+        my_style = mpf.make_mpf_style(base_mpf_style='charles')
+        
+        fig, axlist = mpf.plot(data, type='candle', style=my_style, volume=True, returnfig=True)
+        
+        # Set the title for the figure with padding
+        fig.suptitle(title, y=0.98)  # Adjust the y parameter to move the title down
+        
+        fig.savefig(filename, dpi=300)  # Increased DPI
+        plt.close(fig)
+
+        # Resize image to 3 times its original size
+        image = Image.open(filename)
+        new_size = (image.width * 3, image.height * 3)
+        resized_image = image.resize(new_size, Image.LANCZOS)
+        resized_image.save(filename)
+        logging.debug(f"Resized image with timeframe {timeframe} and ticker {ticker} saved to {filename}")
+    except Exception as e:
+        logging.error(f"Error creating or resizing chart: {e}")
+        raise
+
+# Function to combine two images side by side with increased size
+def combine_images(image1_path, image2_path, output_path='combined_chart.png'):
+    try:
+        logging.debug(f"Combining images {image1_path} and {image2_path} into {output_path}")
+        image1 = Image.open(image1_path)
+        image2 = Image.open(image2_path)
+
+        total_width = image1.width + image2.width
+        max_height = max(image1.height, image2.height)
+
+        combined_image = Image.new('RGB', (total_width, max_height))
+        combined_image.paste(image1, (0, 0))
+        combined_image.paste(image2, (image1.width, 0))
+
+        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
+
+# Function to generate insights
+def generate_insights(image, query, ticker1=None, ticker2=None):
+    try:
+        logging.debug(f"Generating insights for query: {query}")
+
+        # Open and process the image
+        image = Image.open(image).convert('RGB')
+        inputs = processor(text=query, images=image, return_tensors="pt")
+        logging.debug(f"Inputs prepared with shapes {inputs['input_ids'].shape} and {inputs['pixel_values'].shape}")
+        
+        prompt_length = inputs['input_ids'].shape[1]
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        # Generate insights using the model
+        generate_ids = model.generate(**inputs, num_beams=4, max_new_tokens=512)
+        output_text = processor.batch_decode(generate_ids[:, prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+
+        # Replace placeholders with actual ticker names in the insights
+        if ticker1:
+            output_text = output_text.replace("[First Ticker]", ticker1)
+        if ticker2:
+            output_text = output_text.replace("[Second Ticker]", ticker2)
+
+        logging.debug(f"Generated insights: {output_text}")
+
+        return output_text
+    except Exception as e:
+        logging.error(f"Error generating insights: {e}")
+        return f"Error generating insights: {e}"
+
+# Function to handle the Gradio interface
+def gradio_interface(ticker1, start_date, end_date, ticker2, query, analysis_type, interval):
+    try:
+        logging.debug(f"Starting gradio_interface with ticker1: {ticker1}, start_date: {start_date}, end_date: {end_date}, ticker2: {ticker2}, query: {query}, analysis_type: {analysis_type}, interval: {interval}")
+
+        # Fetch data and create charts
+        data1 = fetch_stock_data(ticker1, start=start_date, end=end_date, interval=interval)
+        chart_path1 = '/tmp/chart1.png'
+        create_stock_chart(data1, ticker1, chart_path1, timeframe=interval)
+
+        if analysis_type == 'Comparative Analysis' and ticker2:
+            data2 = fetch_stock_data(ticker2, start=start_date, end=end_date, interval=interval)
+            chart_path2 = '/tmp/chart2.png'
+            create_stock_chart(data2, ticker2, chart_path2, timeframe=interval)
+
+            # Combine the two charts into one image
+            combined_chart_path = combine_images(chart_path1, chart_path2)
+            insights = generate_insights(combined_chart_path, query, ticker1, ticker2)
+            return insights, combined_chart_path
+
+        insights = generate_insights(chart_path1, query, ticker1)
+        return insights, chart_path1
+    except Exception as e:
+        logging.error(f"Error processing image or query: {e}")
+        return f"Error processing image or query: {e}", None
+
+# Button callback functions
+def set_query_trend():
+    return "What are the key trends shown in this chart?"
+
+def set_query_comparative():
+    return "How does [First Ticker] compare to [Second Ticker]?"
+
+def set_query_forecasting():
+    return "Based on the current data, what are the projected trends?"
+
+# Create the Gradio interface
+with gr.Blocks() as interface:
+    gr.Markdown(
+        """
+        # 📈 Price Market Analysis Tool
+        Welcome to the Price Market Analysis Tool! This interface helps you generate insightful analyses of market data. Choose between trend analysis, comparative analysis, and forecasting based on your needs.
+        """
+    )
+
+    with gr.Row():
+        # Input box for first ticker
+        ticker1_input = gr.Textbox(
+            lines=1,
+            placeholder="Enter first ticker (e.g., TSLA)",
+            label="First Ticker",
+        )
+
+        # Input box for second ticker
+        ticker2_input = gr.Textbox(
+            lines=1,
+            placeholder="Enter second ticker for comparative analysis (optional)",
+            label="Second Ticker (Optional)",
+        )
+
+        # Input box for start date
+        start_date_input = gr.Textbox(
+            lines=1,
+            placeholder="Enter start date (e.g., 2023-01-01)",
+            label="Start Date",
+        )
+
+        # Input box for end date
+        end_date_input = gr.Textbox(
+            lines=1,
+            placeholder="Enter end date (e.g., 2024-01-01)",
+            label="End Date",
+        )
+
+        # Input box for text query
+        query_input = gr.Textbox(
+            lines=2,
+            placeholder="Enter your question here...",
+            label="Input Text",
+        )
+
+        # Hidden input for analysis type
+        analysis_type_input = gr.Textbox(
+            lines=1,
+            visible=False,
+            label="Analysis Type"
+        )
+
+        # Dropdown for selecting time frame
+        interval_input = gr.Dropdown(
+            choices=['1m', '5m', '15m', '30m', '60m', '1d', '1wk', '1mo', '3mo'],
+            value='1d',
+            label="Select Time Frame",
+        )
+
+    with gr.Row():
+        trend_button = gr.Button("Trend Analysis")
+        comparative_button = gr.Button("Comparative Analysis")
+        forecasting_button = gr.Button("Forecasting")
+
+    # Output areas for insights and chart
+    output_text = gr.Textbox(
+        lines=10,
+        label="Generated Insights"
+    )
+    output_image = gr.Image(label="Price Chart")
+
+    # Button actions to set query text and analysis type
+    trend_button.click(lambda: ("Trend Analysis", set_query_trend()), outputs=[analysis_type_input, query_input])
+    comparative_button.click(lambda: ("Comparative Analysis", set_query_comparative()), outputs=[analysis_type_input, query_input])
+    forecasting_button.click(lambda: ("Forecasting", set_query_forecasting()), outputs=[analysis_type_input, query_input])
+
+    # Process inputs and generate insights, display chart(s)
+    gr.Interface(gradio_interface,
+                 inputs=[ticker1_input, start_date_input, end_date_input, ticker2_input, query_input, analysis_type_input, interval_input],
+                 outputs=[output_text, output_image]).launch()
+
+# Launch Gradio interface
+if __name__ == "__main__":
+    interface.launch()