Update app.py

app.py

@@ -4,45 +4,36 @@ import matplotlib.pyplot as plt
 import mplfinance as mpf
 from PIL import Image
 import gradio as gr
+import datetime
 import logging
 from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
 import spaces
 
-# Configure logging to write to a file
+# Configure logging
 logging.basicConfig(filename='debug.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Load the ChartGemma model and processor outside the GPU context
+# Load the ChartGemma model and processor
 model = PaliGemmaForConditionalGeneration.from_pretrained("ahmed-masry/chartgemma")
 processor = AutoProcessor.from_pretrained("ahmed-masry/chartgemma")
 
 @spaces.GPU
-def predict(image, input_text, ticker1=None, ticker2=None):
+def predict(image, input_text):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
     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 insights using the model
-    generate_ids = model.generate(**inputs, num_beams=4, max_new_tokens=512)
+    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]
 
-    # 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
 
-# Function to fetch stock data with different intervals
-def fetch_stock_data(ticker='TSLA', start='2023-01-01', end='2024-01-01', interval='1d'):
+def fetch_stock_data(ticker='TSLA', start='2010-01-01', end=None, interval='1d'):
+    if end is None:
+        end = datetime.datetime.now().strftime('%Y-%m-%d')
     try:
         logging.debug(f"Fetching data for {ticker} from {start} to {end} with interval {interval}")
         stock = yf.Ticker(ticker)
@@ -53,7 +44,6 @@ def fetch_stock_data(ticker='TSLA', start='2023-01-01', end='2024-01-01', interv
         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}")
@@ -66,7 +56,6 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d'):
         fig.savefig(filename, dpi=300)
         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)
@@ -76,7 +65,6 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d'):
         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}")
@@ -97,7 +85,6 @@ def combine_images(image1_path, image2_path, output_path='combined_chart.png'):
         logging.error(f"Error combining images: {e}")
         raise
 
-# 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}")
@@ -112,18 +99,16 @@ def gradio_interface(ticker1, start_date, end_date, ticker2, query, analysis_typ
             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 = predict(Image.open(combined_chart_path), query, ticker1, ticker2)
+            insights = predict(Image.open(combined_chart_path), query)
             return insights, combined_chart_path
 
-        insights = predict(Image.open(chart_path1), query, ticker1)
+        insights = predict(Image.open(chart_path1), query)
         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?"
 
@@ -133,7 +118,10 @@ def set_query_comparative():
 def set_query_forecasting():
     return "Based on the current data, what are the projected trends?"
 
-# Create the Gradio interface
+# Default dates
+default_start_date = '2010-01-01'
+default_end_date = datetime.datetime.now().strftime('%Y-%m-%d')
+
 with gr.Blocks() as interface:
     gr.Markdown(
         """
@@ -143,49 +131,44 @@ with gr.Blocks() as interface:
     )
 
     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)",
+            placeholder="Enter start date (e.g., 2010-01-01)",
+            value=default_start_date,
             label="Start Date",
         )
 
-        # Input box for end date
         end_date_input = gr.Textbox(
             lines=1,
-            placeholder="Enter end date (e.g., 2024-01-01)",
+            placeholder=f"Enter end date (e.g., {default_end_date})",
+            value=default_end_date,
             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',
@@ -197,23 +180,19 @@ with gr.Blocks() as interface:
         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 Gradio interface
 if __name__ == "__main__":
     interface.launch()