Update app.py

app.py

@@ -41,68 +41,137 @@ The user provided the additional info about how they would like you to respond:
 total_tokens_used = 0
 
 def format_prompt(message, history):
-    prompt = f"<s>[SYSTEM] {system_instruction} [/SYSTEM]"
+    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
 
-stock_info = {
-    "AAPL": {'name': '애플', 'description': '아이폰을 주력으로 생산하는'},
-    "MSFT": {'name': '마이크로소프트', 'description': '윈도우 운영체제와 오피스 소프트웨어를'},
-    "AMZN": {'name': '아마존', 'description': '전자상거래 및 클라우드 서비스를'},
-    "GOOGL": {'name': '알파벳 (구글)', 'description': '검색 엔진 및 온라인 광고를'},
-    "TSLA": {'name': '테슬라', 'description': '전기자동차와 에너지 저장장치를'}
-}
-
 def get_stock_data(ticker):
-    try:
-        stock = yf.Ticker(ticker)
-        hist = stock.history(period="5d")
-        return hist
-    except Exception as e:
-        return f"데이터를 불러오는 중 오류가 발생했습니다: {e}"
+    stock = yf.Ticker(ticker)
+    hist = stock.history(period="5d")  # 지난 5일간의 주식 데이터를 가져옵니다.
+    return hist
 
 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
+
+    formatted_prompt = format_prompt(prompt, history)
+    output_accumulated = ""
     try:
-        input_tokens = len(tokenizer.encode(prompt))
-        total_tokens_used += input_tokens
-        available_tokens = 32768 - total_tokens_used
-        if available_tokens <= 0:
-            return "Error: 입력이 최대 허용 토큰 수를 초과합니다."
+        # 티커 확인 및 데이터 수집
+        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
 
-        formatted_prompt = format_prompt(prompt, history)
-        ticker = prompt.upper()
-        stock_info_detail = stock_info.get(ticker, None)
-        if stock_info_detail:
-            response_msg = f"{stock_info_detail['name']}은(는) {stock_info_detail['description']} 주력으로 생산하는 기업입니다. 티커는 {ticker}입니다. 원하시는 종목이 맞는가요?"
-            stock_data = get_stock_data(ticker)
-            return f"{response_msg}\n\n---\nTotal tokens used: {total_tokens_used}\nStock Data: {stock_data}"
+            # 추가적인 분석 요청이 있다면, 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:
-            return f"입력하신 '{prompt}'은(는) 지원되는 종목명 또는 티커가 아닙니다. 지원되는 티커: {', '.join(stock_info.keys())}"
+            # 입력이 티커인 경우 처리
+            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) 등입니다. 정확한 종목명 또는 티커를 입력해주세요."
     except Exception as e:
-        return f"Error: {str(e)}\nTotal tokens used: {total_tokens_used}"
+        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,
     show_label=False,
     show_copy_button=True,
-    likeable=True
+    likeable=True,
 )
 
+examples = [
+    ["반드시 한글로 답변할것.", []],
+    ["좋은 종목(티커) 추천해줘", []],
+    ["요약 결론을 제시해", []],
+    ["포트폴리오 분석해줘", []]
+]
+
+css = """
+h1 {
+    font-size: 14px;
+}
+footer {
+    visibility: hidden;
+}
+"""
+
 demo = gr.ChatInterface(
     fn=generate,
     chatbot=mychatbot,
     title="글로벌 자산 분석 및 예측 LLM: BloombAI",
-    css="h1 { font-size: 14px; } footer { visibility: hidden; }",
-    examples=[
-        ["반드시 한글로 답변할것.", []],
-        ["좋은 종목(티커) 추천해줘", []],
-        ["요약 결론을 제시해", []],
-        ["포트폴리오 분석해줘", []]
-    ]
+    retry_btn=None,
+    undo_btn=None,
+    css=css,
+    examples=examples
 )
 
-demo.launch(show_api=False)
+demo.queue().launch(show_api=False)