Update app.py

app.py

@@ -5,14 +5,14 @@ import pandas as pd
 import json
 from datetime import datetime
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import spaces
+from threading import Thread
 
 # 환경 변수 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
 MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
 
-import spaces
-
 class ModelManager:
     def __init__(self):
         self.tokenizer = None
@@ -142,9 +142,6 @@ class ChatHistory:
 chat_history = ChatHistory()
 model_manager = ModelManager()
 
-def get_client():
-    return InferenceClient(MODEL_ID, token=HF_TOKEN)
-
 def analyze_file_content(content, file_type):
     """Analyze file content and return structural summary"""
     if file_type in ['parquet', 'csv']:
@@ -262,17 +259,15 @@ def chat(message, history, uploaded_file, system_message="", max_tokens=4000, te
         
         messages.append({"role": "user", "content": message})
 
-        client = get_client()
         partial_message = ""
         
-        for msg in client.chat_completion(
+        for response in model_manager.generate_response(
             messages,
             max_tokens=max_tokens,
-            stream=True,
             temperature=temperature,
-            top_p=top_p,
+            top_p=top_p
         ):
-            token = msg.choices[0].delta.get('content', None)
+            token = response.choices[0].delta.get('content', '')
             if token:
                 partial_message += token
                 current_history = history + [[message, partial_message]]