Update app.py

app.py

@@ -2,6 +2,11 @@ from huggingface_hub import InferenceClient
 import gradio as gr
 from transformers import GPT2Tokenizer
 import yfinance as yf
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import talib
+import tech_indicators as ti
 
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
@@ -49,9 +54,37 @@ def format_prompt(message, history):
 
 def get_stock_data(ticker):
     stock = yf.Ticker(ticker)
-    hist = stock.history(period="5d")  # 지난 5일간의 주식 데이터를 가져옵니다.
+    hist = stock.history(period="6mo")  # 지난 6개월간의 데이터를 가져옵니다.
     return hist
 
+def apply_technical_indicators(df):
+    df['SMA'] = talib.SMA(df['Close'], timeperiod=20)
+    df['EMA'] = talib.EMA(df['Close'], timeperiod=20)
+    df['RSI'] = talib.RSI(df['Close'], timeperiod=14)
+    macd, macdsignal, macdhist = talib.MACD(df['Close'], fastperiod=12, slowperiod=26, signalperiod=9)
+    df['MACD'] = macd
+    df['MACD_signal'] = macdsignal
+    return df
+
+def plot_technical_indicators(df):
+    plt.figure(figsize=(14, 7))
+    plt.subplot(2, 1, 1)
+    plt.plot(df['Close'], label='Close Price')
+    plt.plot(df['SMA'], label='SMA 20')
+    plt.plot(df['EMA'], label='EMA 20')
+    plt.title('Price Chart with SMA and EMA')
+    plt.legend()
+
+    plt.subplot(2, 1, 2)
+    plt.plot(df['RSI'], label='RSI')
+    plt.title('RSI Chart')
+    plt.legend()
+
+    plt.tight_layout()
+    plt.savefig('/mnt/data/Technical_Indicators.png')
+    plt.close()
+    return '/mnt/data/Technical_Indicators.png'
+
 def generate(prompt, history=[], temperature=0.1, max_new_tokens=10000, top_p=0.95, repetition_penalty=1.0):
     global total_tokens_used
     input_tokens = len(tokenizer.encode(prompt))
@@ -64,82 +97,17 @@ def generate(prompt, history=[], temperature=0.1, max_new_tokens=10000, top_p=0.
     formatted_prompt = format_prompt(prompt, history)
     output_accumulated = ""
     try:
-        # 티커 확인 및 데이터 수집
-        stock_info = get_stock_info(prompt)  # 종목명을 토대로 티커 정보와 기업 설명을 가져옵니다.
-        if stock_info['ticker']:
-            response_msg = f"{stock_info['name']}은(는) {stock_info['description']} 주력으로 생산하는 기업입니다. {stock_info['name']}의 티커는 {stock_info['ticker']}입니다. 원하시는 종목이 맞는가요?"
-            output_accumulated += response_msg
-            yield output_accumulated
-
-            # 추가적인 분석 요청이 있다면, yfinance로 데이터 수집 및 분석
-            stock_data = get_stock_data(stock_info['ticker'])  # 티커를 이용해 주식 데이터를 가져옵니다.
-            stream = client.text_generation(
-                formatted_prompt,
-                temperature=temperature,
-                max_new_tokens=min(max_new_tokens, available_tokens),
-                top_p=top_p,
-                repetition_penalty=repetition_penalty,
-                do_sample=True,
-                seed=42,
-                stream=True
-            )
-            for response in stream:
-                output_part = response['generated_text'] if 'generated_text' in response else str(response)
-                output_accumulated += output_part
-                yield output_accumulated + f"\n\n---\nTotal tokens used: {total_tokens_used}\nStock Data: {stock_data}"
+        ticker = prompt.upper()
+        stock_data = get_stock_data(ticker)
+        if not stock_data.empty:
+            enhanced_data = apply_technical_indicators(stock_data)
+            image_path = plot_technical_indicators(enhanced_data)
+            yield f"Technical analysis for {ticker} completed. See the chart here: {image_path}\n\n---\nTotal tokens used: {total_tokens_used}"
         else:
-            # 입력이 티커인 경우 처리
-            ticker = prompt.upper()
-            if ticker in ['AAPL', 'MSFT', 'AMZN', 'GOOGL', 'TSLA']:
-                stock_info = get_stock_info_by_ticker(ticker)
-                response_msg = f"{stock_info['name']}은(는) {stock_info['description']} 주력으로 생산하는 기업입니다. {stock_info['name']}의 티커는 {stock_info['ticker']}입니다. 원하시는 종목이 맞는가요?"
-                output_accumulated += response_msg
-                yield output_accumulated
-
-                # 추가적인 분석 요청이 있다면, yfinance로 데이터 수집 및 분석
-                stock_data = get_stock_data(stock_info['ticker'])  # 티커를 이용해 주식 데이터를 가져옵니다.
-                stream = client.text_generation(
-                    formatted_prompt,
-                    temperature=temperature,
-                    max_new_tokens=min(max_new_tokens, available_tokens),
-                    top_p=top_p,
-                    repetition_penalty=repetition_penalty,
-                    do_sample=True,
-                    seed=42,
-                    stream=True
-                )
-                for response in stream:
-                    output_part = response['generated_text'] if 'generated_text' in response else str(response)
-                    output_accumulated += output_part
-                    yield output_accumulated + f"\n\n---\nTotal tokens used: {total_tokens_used}\nStock Data: {stock_data}"
-            else:
-                yield f"입력하신 '{prompt}'은(는) 지원되는 종목명 또는 티커가 아닙니다. 현재 지원되는 종목은 애플(AAPL), 마이크로소프트(MSFT), 아마존(AMZN), 알파벳(GOOGL), 테슬라(TSLA) 등입니다. 정확한 종목명 또는 티커를 입력해주세요."
+            yield f"No data available for {ticker}. Please check the ticker and try again."
     except Exception as e:
         yield f"Error: {str(e)}\nTotal tokens used: {total_tokens_used}"
 
-# 티커를 토대로 종목 정보를 제공하는 함수
-def get_stock_info_by_ticker(ticker):
-    stock_info = {
-        "AAPL": {'ticker': 'AAPL', 'name': '애플', 'description': '아이폰을'},
-        "MSFT": {'ticker': 'MSFT', 'name': '마이크로소프트', 'description': '윈도우 운영체제와 오피스 소프트웨어를'},
-        "AMZN": {'ticker': 'AMZN', 'name': '아마존', 'description': '전자상거래 및 클라우드 서비스를'},
-        "GOOGL": {'ticker': 'GOOGL', 'name': '알파벳', 'description': '검색 엔진 및 온라인 광고를'},
-        "TSLA": {'ticker': 'TSLA', 'name': '테슬라', 'description': '전기자동차와 에너지 저장장치를'},
-    }
-    return stock_info.get(ticker, {'ticker': None, 'name': None, 'description': ''})
-
-# 종목명을 토대로 티커와 기업 정보를 제공하는 함수
-def get_stock_info(name):
-    stock_info = {
-        "apple": {'ticker': 'AAPL', 'name': '애플', 'description': '아이폰을'},
-        "microsoft": {'ticker': 'MSFT', 'name': '마이크로소프트', 'description': '윈도우 운영체제와 오피스 소프트웨어를'},
-        "amazon": {'ticker': 'AMZN', 'name': '아마존', 'description': '전자상거래 및 클라우드 서비스를'},
-        "google": {'ticker': 'GOOGL', 'name': '알파벳 (구글)', 'description': '검색 엔진 및 온라인 광고를'},
-        "tesla": {'ticker': 'TSLA', 'name': '테슬라', 'description': '전기자동차와 에너지 저장장치를'},
-        # 추가적인 종목에 대한 정보를 이곳에 구현할 수 있습니다.
-    }
-    return stock_info.get(name.lower(), {'ticker': None, 'name': name, 'description': ''})
-    
 mychatbot = gr.Chatbot(
     avatar_images=["./user.png", "./botm.png"],
     bubble_full_width=False,