Update app.py

app.py

@@ -1,78 +1,71 @@
 import os
 from dotenv import load_dotenv
 import gradio as gr
-from huggingface_hub import InferenceClient
 import pandas as pd
 import json
 from datetime import datetime
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # 환경 변수 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
 MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
 
-from transformers import pipeline
-
 class ModelManager:
     def __init__(self):
-        self.pipe = None
-        self.setup_pipeline()
+        self.tokenizer = None
+        self.model = None
+        self.setup_model()
     
-    def setup_pipeline(self):
+    def setup_model(self):
         try:
-            print("파이프라인 초기화 시작...")
-            self.pipe = pipeline(
-                "text-generation",
-                model=MODEL_ID,
+            print("토크나이저 로딩 시작...")
+            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+            print("토크나이저 로딩 완료")
+            
+            print("모델 로딩 시작...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
                 token=HF_TOKEN,
-                device_map="auto",
-                torch_dtype=torch.float16
+                torch_dtype=torch.float16,
+                device_map="auto"
             )
-            print("파이프라인 초기화 완료")
+            print("모델 로딩 완료")
         except Exception as e:
-            print(f"파이프라인 초기화 중 오류 발생: {e}")
-            raise Exception(f"파이프라인 초기화 실패: {e}")
+            print(f"모델 로딩 중 오류 발생: {e}")
+            raise Exception(f"모델 로딩 실패: {e}")
 
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
         try:
-            # 메시지 형식 변환
-            prompt = ""
-            for msg in messages:
-                role = msg["role"]
-                content = msg["content"]
-                if role == "system":
-                    prompt += f"System: {content}\n"
-                elif role == "user":
-                    prompt += f"User: {content}\n"
-                elif role == "assistant":
-                    prompt += f"Assistant: {content}\n"
-            
-            # 응답 생성
-            response = self.pipe(
-                prompt,
+            # 채팅 템플릿 적용
+            input_ids = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_tensors="pt"
+            ).to(self.model.device)
+
+            # 토큰 생성
+            gen_tokens = self.model.generate(
+                input_ids,
                 max_new_tokens=max_tokens,
+                do_sample=True,
                 temperature=temperature,
                 top_p=top_p,
-                do_sample=True,
-                num_return_sequences=1,
-                pad_token_id=self.pipe.tokenizer.eos_token_id
+                pad_token_id=self.tokenizer.eos_token_id,
+                streamer=TextIteratorStreamer(self.tokenizer, skip_special_tokens=True)
             )
-            
-            # 응답 텍스트 추출 및 스트리밍 시뮬레이션
-            generated_text = response[0]['generated_text'][len(prompt):].strip()
-            words = generated_text.split()
-            
-            # 단어 단위로 스트리밍
-            partial_response = ""
-            for word in words:
-                partial_response += word + " "
+
+            # 응답 디코딩 및 스트리밍
+            response_text = ""
+            for new_text in self.tokenizer.decode(gen_tokens[0], skip_special_tokens=True):
+                response_text += new_text
                 yield type('Response', (), {
                     'choices': [type('Choice', (), {
-                        'delta': {'content': word + " "}
+                        'delta': {'content': new_text}
                     })()]
                 })()
-                
+
         except Exception as e:
             raise Exception(f"응답 생성 실패: {e}")