Update app.py

app.py

@@ -20,15 +20,23 @@ class ModelManager:
     def setup_model(self):
         try:
             print("토크나이저 로딩 시작...")
-            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN,
+                use_fast=True
+            )
             print("토크나이저 로딩 완료")
             
             print("모델 로딩 시작...")
+            # ZERO GPU 설정
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
                 token=HF_TOKEN,
                 torch_dtype=torch.float16,
-                device_map="auto"
+                device_map="balanced",  # ZERO GPU를 위한 balanced 설정
+                max_memory={0: "8GiB"},  # ZERO GPU 메모리 제한
+                offload_folder="offload",  # 오프로드 설정
+                low_cpu_mem_usage=True
             )
             print("모델 로딩 완료")
         except Exception as e:
@@ -37,7 +45,6 @@ class ModelManager:
 
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
         try:
-            # 채팅 템플릿 적용
             input_ids = self.tokenizer.apply_chat_template(
                 messages,
                 tokenize=True,
@@ -45,7 +52,7 @@ class ModelManager:
                 return_tensors="pt"
             ).to(self.model.device)
 
-            # 토큰 생성
+            # ZERO GPU에 최적화된 생성 설정
             gen_tokens = self.model.generate(
                 input_ids,
                 max_new_tokens=max_tokens,
@@ -53,16 +60,18 @@ class ModelManager:
                 temperature=temperature,
                 top_p=top_p,
                 pad_token_id=self.tokenizer.eos_token_id,
-                streamer=TextIteratorStreamer(self.tokenizer, skip_special_tokens=True)
+                use_cache=True,  # 캐시 사용으로 메모리 효율화
+                num_beams=1  # 빔 서치 비활성화로 메모리 절약
             )
 
-            # 응답 디코딩 및 스트리밍
-            response_text = ""
-            for new_text in self.tokenizer.decode(gen_tokens[0], skip_special_tokens=True):
-                response_text += new_text
+            response_text = self.tokenizer.decode(gen_tokens[0][input_ids.shape[1]:], skip_special_tokens=True)
+            
+            # 단어 단위 스트리밍
+            words = response_text.split()
+            for word in words:
                 yield type('Response', (), {
                     'choices': [type('Choice', (), {
-                        'delta': {'content': new_text}
+                        'delta': {'content': word + " "}
                     })()]
                 })()