Update app.py

app.py

@@ -70,28 +70,35 @@ class ModelManager:
                     prompt += f"Assistant: {content}\n"
             prompt += "Assistant: "
 
-            # 토크나이징
-            input_ids = self.tokenizer(
+            # 토크나이징 및 device 설정
+            inputs = 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
             )
+            
+            # 모든 텐서를 GPU로 이동
+            input_ids = inputs.input_ids.to(self.model.device)
+            attention_mask = inputs.attention_mask.to(self.model.device)
 
-            # 디코딩
+            # 생성
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    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로 이동
+            outputs = outputs.cpu()
             generated_text = self.tokenizer.decode(
                 outputs[0][input_ids.shape[1]:],
                 skip_special_tokens=True