Update app.py

app.py

@@ -3,9 +3,9 @@ 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
 
+
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
@@ -43,75 +43,64 @@ The user provided the additional info about how they would like you to respond:
 
 total_tokens_used = 0
 
-def format_prompt(message, history):
-    prompt = "<s>[SYSTEM] {} [/SYSTEM]".format(system_instruction)
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]{bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
-
 def get_stock_data(ticker):
     stock = yf.Ticker(ticker)
-    hist = stock.history(period="6mo")  # 지난 6개월간의 데이터를 가져옵니다.
+    hist = stock.history(period="6mo")  # 지난 6개월간의 주식 데이터를 가져옵니다.
     return hist
 
+def apply_technical_indicators(df):
+    df['SMA'] = df['Close'].rolling(window=20).mean()
+    return df
 
+def plot_technical_indicators(df, ticker):
+    plt.figure(figsize=(10, 5))
+    plt.plot(df['Close'], label='Close Price')
+    plt.plot(df['SMA'], label='20-day SMA')
+    plt.title(f'Technical Indicators for {ticker}')
+    plt.legend()
+    plt.savefig('/mnt/data/Technical_Indicators.png')
+    plt.close()
+    return '/mnt/data/Technical_Indicators.png'
 
+# Gradio 인터페이스 생성을 위한 함수
 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))
-    total_tokens_used += input_tokens
-    available_tokens = 32768 - total_tokens_used
-    if available_tokens <= 0:
-        yield f"Error: 입력이 최대 허용 토큰 수를 초과합니다. Total tokens used: {total_tokens_used}"
-        return
+    if input_tokens >= 32768:
+        return "Error: 입력이 최대 허용 토큰 수를 초과합니다."
 
-    formatted_prompt = format_prompt(prompt, history)
-    output_accumulated = ""
     try:
         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}"
+            stock_data = apply_technical_indicators(stock_data)
+            image_path = plot_technical_indicators(stock_data, ticker)
+            return f"Technical analysis completed. See the chart here: {image_path}"
         else:
-            yield f"No data available for {ticker}. Please check the ticker and try again."
+            return 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}"
-
-mychatbot = gr.Chatbot(
-    avatar_images=["./user.png", "./botm.png"],
-    bubble_full_width=False,
-    show_label=False,
-    show_copy_button=True,
-    likeable=True,
-)
+        return f"Error: {str(e)}"
 
+# Gradio 인터페이스 설정
 examples = [
-    ["반드시 한글로 답변할것.", []],
-    ["좋은 종목(티커) 추천해줘", []],
-    ["요약 결론을 제시해", []],
-    ["포트폴리오 분석해줘", []]
+    ["AAPL", []],
+    ["MSFT", []],
+    ["AMZN", []],
+    ["GOOGL", []],
+    ["TSLA", []]
 ]
 
 css = """
-h1 {
-    font-size: 14px;
-}
-footer {
-    visibility: hidden;
-}
+h1 { font-size: 14px; }
+footer { visibility: hidden; }
 """
 
-demo = gr.ChatInterface(
+demo = gr.Interface(
     fn=generate,
-    chatbot=mychatbot,
-    title="글로벌 자산 분석 및 예측 LLM: BloombAI",
-    retry_btn=None,
-    undo_btn=None,
+    inputs="text",
+    outputs="text",
+    examples=examples,
     css=css,
-    examples=examples
+    title="글로벌 자산 분석 및 예측 LLM: BloombAI"
 )
 
-demo.queue().launch(show_api=False)
+demo.launch()