Update app.py

app.py

@@ -6,7 +6,7 @@ from PIL import Image
 import gradio as gr
 import datetime
 import logging
-from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForPreTraining
 import spaces
 import pandas as pd
 
@@ -14,8 +14,8 @@ import pandas as pd
 logging.basicConfig(filename='debug.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Load the ChartGemma model and processor
-model = PaliGemmaForConditionalGeneration.from_pretrained("ahmed-masry/chartgemma")
-processor = AutoProcessor.from_pretrained("ahmed-masry/chartgemma")
+processor = AutoProcessor.from_pretrained("mobenta/chart_analysis")
+model = AutoModelForPreTraining.from_pretrained("mobenta/chart_analysis")
 
 @spaces.GPU
 def predict(image, input_text):
@@ -59,11 +59,6 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
         # Calculate indicators if selected
         addplot = []
         if indicators:
-            if 'VWAP' in indicators:
-                vwap = (data['Close'] * data['Volume']).cumsum() / data['Volume'].cumsum()
-                addplot.append(mpf.make_addplot(vwap, color='purple'))
-            if 'Volume' in indicators:
-                addplot.append(mpf.make_addplot(data['Volume'], panel=1, type='bar', color='g', ylabel='Volume'))
             if 'RSI' in indicators:
                 delta = data['Close'].diff(1)
                 gain = delta.where(delta > 0, 0)
@@ -73,53 +68,6 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
                 rs = avg_gain / avg_loss
                 rsi = 100 - (100 / (1 + rs))
                 addplot.append(mpf.make_addplot(rsi, panel=2, color='orange', ylabel='RSI'))
-            if 'Ichimoku Cloud' in indicators:
-                logging.debug("Calculating Ichimoku Cloud")
-                nine_period_high = data['High'].rolling(window=9).max()
-                nine_period_low = data['Low'].rolling(window=9).min()
-                data['tenkan_sen'] = (nine_period_high + nine_period_low) / 2
-
-                period26_high = data['High'].rolling(window=26).max()
-                period26_low = data['Low'].rolling(window=26).min()
-                data['kijun_sen'] = (period26_high + period26_low) / 2
-
-                data['senkou_span_a'] = ((data['tenkan_sen'] + data['kijun_sen']) / 2).shift(26)
-
-                period52_high = data['High'].rolling(window=52).max()
-                period52_low = data['Low'].rolling(window=52).min()
-                data['senkou_span_b'] = ((period52_high + period52_low) / 2).shift(26)
-
-                data['chikou_span'] = data['Close'].shift(-26)
-
-                addplot.append(mpf.make_addplot(data['tenkan_sen'], color='red'))
-                addplot.append(mpf.make_addplot(data['kijun_sen'], color='blue'))
-                addplot.append(mpf.make_addplot(data['senkou_span_a'], color='green'))
-                addplot.append(mpf.make_addplot(data['senkou_span_b'], color='brown'))
-                addplot.append(mpf.make_addplot(data['chikou_span'], color='purple'))
-            if 'Bollinger Bands' in indicators:
-                logging.debug("Calculating Bollinger Bands")
-                rolling_mean = data['Close'].rolling(window=20).mean()
-                rolling_std = data['Close'].rolling(window=20).std()
-                data['upper_band'] = rolling_mean + (rolling_std * 2)
-                data['lower_band'] = rolling_mean - (rolling_std * 2)
-                addplot.append(mpf.make_addplot(data['upper_band'], color='blue'))
-                addplot.append(mpf.make_addplot(data['lower_band'], color='blue'))
-            if 'Pivot Levels' in indicators:
-                logging.debug("Calculating Pivot Levels")
-                data['pivot'] = (data['High'] + data['Low'] + data['Close']) / 3
-                data['r1'] = (2 * data['pivot']) - data['Low']
-                data['s1'] = (2 * data['pivot']) - data['High']
-                data['r2'] = data['pivot'] + (data['High'] - data['Low'])
-                data['s2'] = data['pivot'] - (data['High'] - data['Low'])
-                data['r3'] = data['High'] + 2 * (data['pivot'] - data['Low'])
-                data['s3'] = data['Low'] - 2 * (data['High'] - data['pivot'])
-                addplot.append(mpf.make_addplot(data['pivot'], color='black'))
-                addplot.append(mpf.make_addplot(data['r1'], color='green'))
-                addplot.append(mpf.make_addplot(data['s1'], color='red'))
-                addplot.append(mpf.make_addplot(data['r2'], color='green', linestyle='dashed'))
-                addplot.append(mpf.make_addplot(data['s2'], color='red', linestyle='dashed'))
-                addplot.append(mpf.make_addplot(data['r3'], color='green', linestyle='dotted'))
-                addplot.append(mpf.make_addplot(data['s3'], color='red', linestyle='dotted'))
             if 'SMA21' in indicators:
                 logging.debug("Calculating SMA 21")
                 sma_21 = data['Close'].rolling(window=21).mean()
@@ -132,6 +80,18 @@ def create_stock_chart(data, ticker, filename='chart.png', timeframe='1d', indic
                 logging.debug("Calculating SMA 200")
                 sma_200 = data['Close'].rolling(window=200).mean()
                 addplot.append(mpf.make_addplot(sma_200, color='brown', linestyle='dashed'))
+            if 'VWAP' in indicators:
+                logging.debug("Calculating VWAP")
+                vwap = (data['Volume'] * (data['High'] + data['Low'] + data['Close']) / 3).cumsum() / data['Volume'].cumsum()
+                addplot.append(mpf.make_addplot(vwap, color='blue', linestyle='dashed'))
+            if 'Bollinger Bands' in indicators:
+                logging.debug("Calculating Bollinger Bands")
+                sma = data['Close'].rolling(window=20).mean()
+                std = data['Close'].rolling(window=20).std()
+                upper_band = sma + (std * 2)
+                lower_band = sma - (std * 2)
+                addplot.append(mpf.make_addplot(upper_band, color='green', linestyle='dashed'))
+                addplot.append(mpf.make_addplot(lower_band, color='green', linestyle='dashed'))
 
         fig, axlist = mpf.plot(data, type='candle', style=my_style, volume=True, addplot=addplot, returnfig=True)
         fig.suptitle(title, y=0.98)
@@ -224,6 +184,7 @@ with gr.Blocks() as interface:
             label="Second Ticker (Optional)",
         )
 
+    with gr.Row():
         start_date_input = gr.Textbox(
             lines=1,
             placeholder="Enter start date (e.g., 2010-01-01)",
@@ -233,53 +194,107 @@ with gr.Blocks() as interface:
 
         end_date_input = gr.Textbox(
             lines=1,
-            placeholder=f"Enter end date (e.g., {default_end_date})",
+            placeholder=f"Enter end date (default: {default_end_date})",
             value=default_end_date,
             label="End Date",
         )
 
-        query_input = gr.Textbox(
-            lines=2,
-            placeholder="Enter your question here...",
+    with gr.Row():
+        input_text = gr.Textbox(
+            lines=3,
+            placeholder="Enter your input text",
             label="Input Text",
         )
 
         analysis_type_input = gr.Textbox(
             lines=1,
+            placeholder="Analysis Type",
+            label="",
+            visible=False,
+        )
+
+        query_input = gr.Textbox(
+            lines=3,
+            placeholder="Query",
+            label="",
             visible=False,
-            label="Analysis Type"
         )
 
+    with gr.Row():
         interval_input = gr.Dropdown(
-            choices=['1m', '5m', '15m', '30m', '60m', '1d', '1wk', '1mo', '3mo'],
-            value='1d',
+            choices=["1d", "1wk", "1mo"],
+            value="1d",
             label="Select Time Frame",
         )
 
+    with gr.Row():
         indicator_input = gr.CheckboxGroup(
-            choices=['VWAP', 'Volume', 'RSI', 'Ichimoku Cloud', 'Bollinger Bands', 'Pivot Levels', 'SMA21', 'SMA50', 'SMA200'],
+            choices=["VWAP", "Volume", "RSI", "Ichimoku Cloud", "Bollinger Bands", "Pivot Levels", "SMA21", "SMA50", "SMA200"],
             label="Select Indicators",
-            value=['Volume']  # Default to show Volume
         )
 
     with gr.Row():
         trend_button = gr.Button("Trend Analysis")
         comparative_button = gr.Button("Comparative Analysis")
         forecasting_button = gr.Button("Forecasting")
+        submit_button = gr.Button("Submit")
+        clear_button = gr.Button("Clear")
+
+    output_text = gr.Textbox(lines=5, label="Generated Insights")
+    output_image = gr.Image(type="filepath", label="Price Chart")
+
+    trend_button.click(
+        fn=lambda: "Trend Analysis",
+        inputs=[],
+        outputs=[analysis_type_input],
+    ).then(
+        set_query_trend,
+        inputs=[],
+        outputs=[query_input],
+    ).then(
+        gradio_interface,
+        inputs=[ticker1_input, start_date_input, end_date_input, ticker2_input, query_input, analysis_type_input, interval_input, indicator_input],
+        outputs=[output_text, output_image],
+    )
+
+    comparative_button.click(
+        fn=lambda: "Comparative Analysis",
+        inputs=[],
+        outputs=[analysis_type_input],
+    ).then(
+        set_query_comparative,
+        inputs=[],
+        outputs=[query_input],
+    ).then(
+        gradio_interface,
+        inputs=[ticker1_input, start_date_input, end_date_input, ticker2_input, query_input, analysis_type_input, interval_input, indicator_input],
+        outputs=[output_text, output_image],
+    )
 
-    output_text = gr.Textbox(
-        lines=10,
-        label="Generated Insights"
+    forecasting_button.click(
+        fn=lambda: "Forecasting",
+        inputs=[],
+        outputs=[analysis_type_input],
+    ).then(
+        set_query_forecasting,
+        inputs=[],
+        outputs=[query_input],
+    ).then(
+        gradio_interface,
+        inputs=[ticker1_input, start_date_input, end_date_input, ticker2_input, query_input, analysis_type_input, interval_input, indicator_input],
+        outputs=[output_text, output_image],
     )
-    output_image = gr.Image(label="Price Chart")
 
-    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])
+    submit_button.click(
+        gradio_interface,
+        inputs=[ticker1_input, start_date_input, end_date_input, ticker2_input, input_text, analysis_type_input, interval_input, indicator_input],
+        outputs=[output_text, output_image],
+    )
 
-    gr.Interface(gradio_interface,
-                 inputs=[ticker1_input, start_date_input, end_date_input, ticker2_input, query_input, analysis_type_input, interval_input, indicator_input],
-                 outputs=[output_text, output_image])
+    clear_button.click(
+        fn=lambda: ("", "", "", "", "", "", "", []),
+        inputs=[],
+        outputs=[ticker1_input, start_date_input, end_date_input, ticker2_input, input_text, analysis_type_input, interval_input, indicator_input],
+    )
 
-if __name__ == "__main__":
-    interface.launch()
+interface.launch(debug=True)