Update app.py

app.py

@@ -22,6 +22,7 @@ class ModelManager:
         if self.model is None or self.tokenizer is None:
             self.setup_model()
 
+    @spaces.GPU
     def setup_model(self):
         try:
             print("토크나이저 로딩 시작...")
@@ -36,12 +37,11 @@ class ModelManager:
             print("토크나이저 로딩 완료")
             
             print("모델 로딩 시작...")
-            # CUDA 초기화 방지를 위한 설정
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
                 token=HF_TOKEN,
                 torch_dtype=torch.float16,
-                device_map=None,  # 초기에는 device_map을 설정하지 않음
+                device_map="auto",
                 trust_remote_code=True,
                 low_cpu_mem_usage=True
             )
@@ -52,41 +52,11 @@ class ModelManager:
             raise Exception(f"모델 로딩 실패: {e}")
 
     @spaces.GPU
-    def generate_text(self, prompt, max_tokens, temperature, top_p):
-        try:
-            # GPU 컨텍스트 내에서 device 설정
-            self.model = self.model.to("cuda")
-            input_ids = self.tokenizer.encode(
-                prompt,
-                return_tensors="pt",
-                add_special_tokens=True
-            ).to("cuda")
-
-            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
-                )
-
-            # CPU로 다시 이동
-            self.model = self.model.to("cpu")
-            return self.tokenizer.decode(
-                output_ids[0][input_ids.shape[1]:],
-                skip_special_tokens=True
-            )
-        except Exception as e:
-            if self.model.device.type == "cuda":
-                self.model = self.model.to("cpu")
-            raise Exception(f"텍스트 생성 실패: {e}")
-
     def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
         try:
+            # 모델이 로드되어 있는지 확인
+            self.ensure_model_loaded()
+
             # 입력 텍스트 준비
             prompt = ""
             for msg in messages:
@@ -100,8 +70,32 @@ class ModelManager:
                     prompt += f"Assistant: {content}\n"
             prompt += "Assistant: "
 
-            # spaces.GPU 데코레이터가 적용된 메서드 호출
-            generated_text = self.generate_text(prompt, max_tokens, temperature, top_p)
+            # 토크나이징
+            input_ids = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=4096
+            ).input_ids
+
+            # 생성
+            outputs = 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(
+                outputs[0][input_ids.shape[1]:],
+                skip_special_tokens=True
+            )
 
             # 단어 단위로 스트리밍
             words = generated_text.split()