Update app.py

app.py

@@ -25,10 +25,11 @@ class ModelManager:
             print("토크나이저 로딩 시작...")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 MODEL_ID,
+                use_fast=True,
                 token=HF_TOKEN,
                 trust_remote_code=True
             )
-            if self.tokenizer.pad_token is None:
+            if not self.tokenizer.pad_token:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             print("토크나이저 로딩 완료")
             
@@ -38,9 +39,16 @@ class ModelManager:
                 token=HF_TOKEN,
                 torch_dtype=torch.bfloat16,
                 device_map="auto",
-                trust_remote_code=True
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
             )
+            self.model.eval()  # 평가 모드로 설정
             print("모델 로딩 완료")
+            
+            # 모델과 토크나이저가 제대로 로드되었는지 확인
+            if self.model is None or self.tokenizer is None:
+                raise Exception("모델 또는 토크나이저가 제대로 초기화되지 않았습니다.")
+                
         except Exception as e:
             print(f"모델 로딩 중 오류 발생: {e}")
             raise Exception(f"모델 로딩 실패: {e}")
@@ -48,59 +56,54 @@ class ModelManager:
     @spaces.GPU
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
         try:
-            # 메시지 포맷팅
-            formatted_messages = []
+            if self.model is None or self.tokenizer is None:
+                raise Exception("모델이 초기화되지 않았습니다.")
+
+            # 입력 텍스트 준비
+            prompt = ""
             for msg in messages:
-                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(
+                role = msg["role"]
+                content = msg["content"]
+                if role == "system":
+                    prompt += f"System: {content}\n"
+                elif role == "user":
+                    prompt += f"Human: {content}\n"
+                elif role == "assistant":
+                    prompt += f"Assistant: {content}\n"
+            prompt += "Assistant: "  # 응답 시작 프롬프트
+
+            # 입력 인코딩
+            input_ids = self.tokenizer.encode(
                 prompt,
                 return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=4096
+                add_special_tokens=True
             ).to(self.model.device)
-            
-            # 스트리머 설정
-            streamer = TextIteratorStreamer(
-                self.tokenizer,
-                timeout=10.,
-                skip_prompt=True,
+
+            # 응답 생성
+            with torch.no_grad():
+                output_ids = self.model.generate(
+                    input_ids,
+                    max_new_tokens=max_tokens,
+                    do_sample=True,
+                    temperature=temperature,
+                    top_p=top_p,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    num_return_sequences=1
+                )
+
+            # 응답 디코딩 및 스트리밍
+            generated_text = self.tokenizer.decode(
+                output_ids[0][input_ids.shape[1]:],
                 skip_special_tokens=True
             )
 
-            # 생성 설정
-            generate_kwargs = dict(
-                **inputs,
-                streamer=streamer,
-                max_new_tokens=max_tokens,
-                do_sample=True,
-                temperature=temperature,
-                top_p=top_p,
-                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
+            # 단어 단위로 스트리밍
+            words = generated_text.split()
+            for word in words:
                 yield type('Response', (), {
                     'choices': [type('Choice', (), {
-                        'delta': {'content': new_text}
+                        'delta': {'content': word + " "}
                     })()]
                 })()