Update app.py

app.py

@@ -66,16 +66,33 @@ def generate(prompt, history=[], temperature=0.1, max_new_tokens=10000, top_p=0.
     formatted_prompt = format_prompt(prompt, history)
     output_accumulated = ""
     try:
-        stock_data = get_stock_data("AAPL")  # 예시로 'AAPL' 티커 데이터를 가져옵니다.
-        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}"
+        # 티커 확인 및 데이터 수집
+        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}"
+        else:
+            yield "입력하신 종목명을 확인할 수 없습니다. 정확한 종목명을 입력해주세요."
     except Exception as e:
         yield f"Error: {str(e)}\nTotal tokens used: {total_tokens_used}"
 
+# 종목명을 토대로 티커와 기업 정보를 제공하는 함수
+def get_stock_info(name):
+    if name.lower() == "apple":
+        return {'ticker': 'AAPL', 'name': '애플', 'description': '아이폰을'}
+    # 추가적인 종목에 대한 정보를 이곳에 구현할 수 있습니다.
+    return {'ticker': None, 'name': name, 'description': ''}
+
+
 mychatbot = gr.Chatbot(
     avatar_images=["./user.png", "./botm.png"],
     bubble_full_width=False,