Update app.py

app.py

@@ -6,15 +6,36 @@ import pandas as pd
 from typing import List, Tuple
 import json
 from datetime import datetime
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import spaces
 
 # 환경 변수 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
+MODEL_ID = "CohereForAI/c4ai-command-r-plus-08-2024"
 
-# LLM Models Definition
-LLM_MODELS = {
-    "Cohere c4ai-crp-08-2024": "CohereForAI/c4ai-command-r-plus-08-2024",  # Default
-    "Meta Llama3.3-70B": "meta-llama/Llama-3.3-70B-Instruct"    # Backup model
-}
+class ModelManager:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.setup_model()
+    
+    def setup_model(self):
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN,
+                trust_remote_code=True
+            )
+            self.model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True
+            )
+        except Exception as e:
+            raise Exception(f"Model loading failed: {e}")
 
 class ChatHistory:
     def __init__(self):
@@ -34,17 +55,15 @@ class ChatHistory:
         self.save_history()
 
     def format_for_display(self):
-        # Gradio Chatbot 컴포넌트에 맞는 형식으로 변환
         formatted = []
         for conv in self.history:
             formatted.append([
-                conv["messages"][0]["content"],  # user message
-                conv["messages"][1]["content"]   # assistant message
+                conv["messages"][0]["content"],
+                conv["messages"][1]["content"]
             ])
         return formatted
 
     def get_messages_for_api(self):
-        # API 호출을 위한 메시지 형식
         messages = []
         for conv in self.history:
             messages.extend([
@@ -73,15 +92,12 @@ class ChatHistory:
             print(f"히스토리 로드 실패: {e}")
             self.history = []
 
-
-# 전역 ChatHistory 인스턴스 생성
+# 전역 인스턴스 생성
 chat_history = ChatHistory()
+model_manager = ModelManager()
 
-def get_client(model_name="Cohere c4ai-crp-08-2024"):
-    try:
-        return InferenceClient(LLM_MODELS[model_name], token=HF_TOKEN)
-    except Exception:
-        return InferenceClient(LLM_MODELS["Meta Llama3.3-70B"], token=HF_TOKEN)
+def get_client():
+    return InferenceClient(MODEL_ID, token=HF_TOKEN)
 
 def analyze_file_content(content, file_type):
     """Analyze file content and return structural summary"""
@@ -164,12 +180,10 @@ def chat(message, history, uploaded_file, system_message="", max_tokens=4000, te
 3. 🎯 질문의 의도를 정확히 파악하여 맞춤형 답변
 4. 📚 필요한 경우 업로드된 파일 내용을 참고하여 구체적인 도움 제공
 5. ✨ 추가적인 통찰과 제안을 통한 가치 있는 대화
-
 항상 예의 바르고 친절하게 응답하며, 필요한 경우 구체적인 예시나 설명을 추가하여 
 이해를 돕겠습니다."""
 
     try:
-        # 파일 업로드 처리
         if uploaded_file:
             content, file_type = read_uploaded_file(uploaded_file)
             if file_type == "error":
@@ -193,10 +207,8 @@ def chat(message, history, uploaded_file, system_message="", max_tokens=4000, te
 4. ✨ 개선 제안
 5. 💬 추가 질문이나 필요한 설명"""
 
-        # 메시지 처리
         messages = [{"role": "system", "content": system_prefix + system_message}]
         
-        # 이전 대화 히스토리 추가
         if history:
             for user_msg, assistant_msg in history:
                 messages.append({"role": "user", "content": user_msg})
@@ -204,7 +216,6 @@ def chat(message, history, uploaded_file, system_message="", max_tokens=4000, te
         
         messages.append({"role": "user", "content": message})
 
-        # API 호출 및 응답 처리
         client = get_client()
         partial_message = ""
         
@@ -221,7 +232,6 @@ def chat(message, history, uploaded_file, system_message="", max_tokens=4000, te
                 current_history = history + [[message, partial_message]]
                 yield "", current_history
 
-        # 완성된 대화 저장
         chat_history.add_conversation(message, partial_message)
         
     except Exception as e:
@@ -230,18 +240,16 @@ def chat(message, history, uploaded_file, system_message="", max_tokens=4000, te
         yield "", history + [[message, error_msg]]
 
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", title="GiniGEN 🤖") as demo:
-    # 기존 히스토리 로드
     initial_history = chat_history.format_for_display()
     with gr.Row():
         with gr.Column(scale=2):
             chatbot = gr.Chatbot(
-                value=initial_history,  # 저장된 히스토리로 초기화
+                value=initial_history,
                 height=600, 
                 label="대화창 💬",
                 show_label=True
             )    
 
-
             msg = gr.Textbox(
                 label="메시지 입력",
                 show_label=False,
@@ -266,7 +274,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", title="GiniGEN 🤖") as demo
                 temperature = gr.Slider(minimum=0, maximum=1, value=0.7, label="창의성 수준 🌡️")
                 top_p = gr.Slider(minimum=0, maximum=1, value=0.9, label="응답 다양성 📈")
 
-    # 예시 질문
     gr.Examples(
         examples=[
             ["안녕하세요! 어떤 도움이 필요하신가요? 🤝"],
@@ -278,12 +285,10 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", title="GiniGEN 🤖") as demo
         inputs=msg,
     )
 
-    # 대화내용 지우기 버튼에 히스토리 초기화 기능 추가
     def clear_chat():
         chat_history.clear_history()
         return None, None
 
-    # 이벤트 바인딩
     msg.submit(
         chat,
         inputs=[msg, chatbot, file_upload, system_message, max_tokens, temperature, top_p],
@@ -301,7 +306,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", title="GiniGEN 🤖") as demo
         outputs=[msg, chatbot]
     )
 
-    # 파일 업로드시 자동 분석
     file_upload.change(
         lambda: "파일 분석을 시작합니다...",
         outputs=msg
@@ -312,4 +316,4 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", title="GiniGEN 🤖") as demo
     )
 
 if __name__ == "__main__":
-    demo.launch()        
+    demo.launch()