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,64 +43,75 @@ 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))
-    if input_tokens >= 32768:
-        return "Error: 입력이 최대 허용 토큰 수를 초과합니다."
+    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
 
+    formatted_prompt = format_prompt(prompt, history)
+    output_accumulated = ""
     try:
         ticker = prompt.upper()
         stock_data = get_stock_data(ticker)
         if not stock_data.empty:
-            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}"
+            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:
-            return f"No data available for {ticker}. Please check the ticker and try again."
+            yield f"No data available for {ticker}. Please check the ticker and try again."
     except Exception as e:
-        return f"Error: {str(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,
+)
 
-# Gradio 인터페이스 설정
 examples = [
-    ["AAPL", []],
-    ["MSFT", []],
-    ["AMZN", []],
-    ["GOOGL", []],
-    ["TSLA", []]
+    ["반드시 한글로 답변할것.", []],
+    ["좋은 종목(티커) 추천해줘", []],
+    ["요약 결론을 제시해", []],
+    ["포트폴리오 분석해줘", []]
 ]
 
 css = """
-h1 { font-size: 14px; }
-footer { visibility: hidden; }
+h1 {
+    font-size: 14px;
+}
+footer {
+    visibility: hidden;
+}
 """
 
-demo = gr.Interface(
+demo = gr.ChatInterface(
     fn=generate,
-    inputs="text",
-    outputs="text",
-    examples=examples,
+    chatbot=mychatbot,
+    title="글로벌 자산 분석 및 예측 LLM: BloombAI",
+    retry_btn=None,
+    undo_btn=None,
     css=css,
-    title="글로벌 자산 분석 및 예측 LLM: BloombAI"
+    examples=examples
 )
 
-demo.launch()
+demo.queue().launch(show_api=False)