Update app.py

app.py

@@ -17,13 +17,11 @@ class ModelManager:
     def __init__(self):
         self.tokenizer = None
         self.model = None
-        # 초기화는 첫 요청 시에 수행
     
     def ensure_model_loaded(self):
         if self.model is None or self.tokenizer is None:
             self.setup_model()
 
-    @spaces.GPU
     def setup_model(self):
         try:
             print("토크나이저 로딩 시작...")
@@ -41,10 +39,11 @@ class ModelManager:
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
                 token=HF_TOKEN,
-                torch_dtype=torch.bfloat16,
+                torch_dtype=torch.float16,
                 device_map="auto",
                 trust_remote_code=True,
-                low_cpu_mem_usage=True
+                low_cpu_mem_usage=True,
+                max_memory={0: "13GB"}  # GPU 메모리 제한
             )
             self.model.eval()
             print("모델 로딩 완료")
@@ -54,32 +53,14 @@ class ModelManager:
             raise Exception(f"모델 로딩 실패: {e}")
 
     @spaces.GPU
-    def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
+    def generate_text(self, prompt, max_tokens, temperature, top_p):
         try:
-            # 모델이 로드되어 있는지 확인
-            self.ensure_model_loaded()
-
-            # 입력 텍스트 준비
-            prompt = ""
-            for msg in messages:
-                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",
                 add_special_tokens=True
             ).to(self.model.device)
 
-            # 응답 생성
             with torch.no_grad():
                 output_ids = self.model.generate(
                     input_ids,
@@ -92,11 +73,30 @@ class ModelManager:
                     num_return_sequences=1
                 )
 
-            # 응답 디코딩
-            generated_text = self.tokenizer.decode(
+            return self.tokenizer.decode(
                 output_ids[0][input_ids.shape[1]:],
                 skip_special_tokens=True
             )
+        except Exception as e:
+            raise Exception(f"텍스트 생성 실패: {e}")
+
+    def generate_response(self, messages, max_tokens=4000, temperature=0.7, top_p=0.9):
+        try:
+            # 입력 텍스트 준비
+            prompt = ""
+            for msg in messages:
+                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: "
+
+            # spaces.GPU 데코레이터가 적용된 메서드 호출
+            generated_text = self.generate_text(prompt, max_tokens, temperature, top_p)
 
             # 단어 단위로 스트리밍
             words = generated_text.split()