Update app.py

app.py

@@ -11,67 +11,72 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 HF_TOKEN = os.getenv("HF_TOKEN")
 MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
 
+import spaces
+
 class ModelManager:
     def __init__(self):
         self.tokenizer = None
         self.model = None
         self.setup_model()
     
+    @spaces.GPU
     def setup_model(self):
         try:
             print("토크나이저 로딩 시작...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                MODEL_ID,
-                token=HF_TOKEN,
-                use_fast=True
-            )
+            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
             print("토크나이저 로딩 완료")
             
             print("모델 로딩 시작...")
-            # ZERO GPU 설정
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
-                token=HF_TOKEN,
-                torch_dtype=torch.float16,
-                device_map="balanced",  # ZERO GPU를 위한 balanced 설정
-                max_memory={0: "8GiB"},  # ZERO GPU 메모리 제한
-                offload_folder="offload",  # 오프로드 설정
-                low_cpu_mem_usage=True
+                torch_dtype=torch.bfloat16,
+                device_map="auto"
             )
             print("모델 로딩 완료")
         except Exception as e:
             print(f"모델 로딩 중 오류 발생: {e}")
             raise Exception(f"모델 로딩 실패: {e}")
 
+    @spaces.GPU
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
         try:
+            conversation = []
+            for msg in messages:
+                conversation.append({"role": msg["role"], "content": msg["content"]})
+
             input_ids = self.tokenizer.apply_chat_template(
-                messages,
-                tokenize=True,
-                add_generation_prompt=True,
-                return_tensors="pt"
-            ).to(self.model.device)
-
-            # ZERO GPU에 최적화된 생성 설정
-            gen_tokens = self.model.generate(
-                input_ids,
+                conversation,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+            inputs = self.tokenizer(input_ids, return_tensors="pt").to(0)
+            
+            streamer = TextIteratorStreamer(
+                self.tokenizer,
+                timeout=10.,
+                skip_prompt=True,
+                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.eos_token_id,
-                use_cache=True,  # 캐시 사용으로 메모리 효율화
-                num_beams=1  # 빔 서치 비활성화로 메모리 절약
+                eos_token_id=[255001]
             )
-
-            response_text = self.tokenizer.decode(gen_tokens[0][input_ids.shape[1]:], skip_special_tokens=True)
             
-            # 단어 단위 스트리밍
-            words = response_text.split()
-            for word in words:
+            thread = Thread(target=self.model.generate, kwargs=generate_kwargs)
+            thread.start()
+
+            buffer = ""
+            for new_text in streamer:
+                buffer += new_text
                 yield type('Response', (), {
                     'choices': [type('Choice', (), {
-                        'delta': {'content': word + " "}
+                        'delta': {'content': new_text}
                     })()]
                 })()