Update app.py

app.py

@@ -41,137 +41,68 @@ 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)
+    prompt = f"<s>[SYSTEM] {system_instruction} [/SYSTEM]"
     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):
-    stock = yf.Ticker(ticker)
-    hist = stock.history(period="5d")  # 지난 5일간의 주식 데이터를 가져옵니다.
-    return hist
+    try:
+        stock = yf.Ticker(ticker)
+        hist = stock.history(period="5d")
+        return hist
+    except Exception as e:
+        return f"데이터를 불러오는 중 오류가 발생했습니다: {e}"
 
 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:
-        # 티커 확인 및 데이터 수집
-        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
+        input_tokens = len(tokenizer.encode(prompt))
+        total_tokens_used += input_tokens
+        available_tokens = 32768 - total_tokens_used
+        if available_tokens <= 0:
+            return "Error: 입력이 최대 허용 토큰 수를 초과합니다."
 
-            # 추가적인 분석 요청이 있다면, 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}"
+        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}"
         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) 등입니다. 정확한 종목명 또는 티커를 입력해주세요."
+            return f"입력하신 '{prompt}'은(는) 지원되는 종목명 또는 티커가 아닙니다. 지원되는 티커: {', '.join(stock_info.keys())}"
     except Exception as e:
-        yield f"Error: {str(e)}\nTotal tokens used: {total_tokens_used}"
+        return 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,
+    examples=[
+        ["반드시 한글로 답변할것.", []],
+        ["좋은 종목(티커) 추천해줘", []],
+        ["요약 결론을 제시해", []],
+        ["포트폴리오 분석해줘", []]
+    ]
 )
 
-examples = [
-    ["반드시 한글로 답변할것.", []],
-    ["좋은 종목(티커) 추천해줘", []],
-    ["요약 결론을 제시해", []],
-    ["포트폴리오 분석해줘", []]
-]
-
-css = """
-h1 {
-    font-size: 14px;
-}
-footer {
-    visibility: hidden;
-}
-"""
-
 demo = gr.ChatInterface(
     fn=generate,
     chatbot=mychatbot,
     title="글로벌 자산 분석 및 예측 LLM: BloombAI",
-    retry_btn=None,
-    undo_btn=None,
-    css=css,
-    examples=examples
+    css="h1 { font-size: 14px; } footer { visibility: hidden; }"
 )
 
-demo.queue().launch(show_api=False)
+demo.launch(show_api=False)