Update app.py

app.py

@@ -23,14 +23,22 @@ class ModelManager:
     def setup_model(self):
         try:
             print("토크나이저 로딩 시작...")
-            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN,
+                trust_remote_code=True
+            )
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
             print("토크나이저 로딩 완료")
             
             print("모델 로딩 시작...")
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
+                token=HF_TOKEN,
                 torch_dtype=torch.bfloat16,
-                device_map="auto"
+                device_map="auto",
+                trust_remote_code=True
             )
             print("모델 로딩 완료")
         except Exception as e:
@@ -40,17 +48,29 @@ class ModelManager:
     @spaces.GPU
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
         try:
-            conversation = []
+            # 메시지 포맷팅
+            formatted_messages = []
             for msg in messages:
-                conversation.append({"role": msg["role"], "content": msg["content"]})
-
-            input_ids = self.tokenizer.apply_chat_template(
-                conversation,
-                tokenize=False,
-                add_generation_prompt=True
-            )
-            inputs = self.tokenizer(input_ids, return_tensors="pt").to(0)
+                if msg["role"] == "system":
+                    formatted_messages.append(f"System: {msg['content']}\n")
+                elif msg["role"] == "user":
+                    formatted_messages.append(f"User: {msg['content']}\n")
+                elif msg["role"] == "assistant":
+                    formatted_messages.append(f"Assistant: {msg['content']}\n")
+            
+            # 입력 텍스트 생성
+            prompt = "".join(formatted_messages)
+            
+            # 토크나이징
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=4096
+            ).to(self.model.device)
             
+            # 스트리머 설정
             streamer = TextIteratorStreamer(
                 self.tokenizer,
                 timeout=10.,
@@ -58,6 +78,7 @@ class ModelManager:
                 skip_special_tokens=True
             )
 
+            # 생성 설정
             generate_kwargs = dict(
                 **inputs,
                 streamer=streamer,
@@ -65,12 +86,15 @@ class ModelManager:
                 do_sample=True,
                 temperature=temperature,
                 top_p=top_p,
-                eos_token_id=[255001]
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id
             )
             
+            # 비동기 생성
             thread = Thread(target=self.model.generate, kwargs=generate_kwargs)
             thread.start()
 
+            # 응답 스트리밍
             buffer = ""
             for new_text in streamer:
                 buffer += new_text
@@ -81,6 +105,7 @@ class ModelManager:
                 })()
 
         except Exception as e:
+            print(f"응답 생성 중 오류 발생: {e}")
             raise Exception(f"응답 생성 실패: {e}")
 
 class ChatHistory: