Auto commit at 07-2025-08 0:31:54

app.py

@@ -56,8 +56,7 @@ try:
             torch_dtype=torch.float16,
             trust_remote_code=True,
             device_map=None,
-            low_cpu_mem_usage=True,
-            # max_memory={0: "4GB"}  # GPU 메모리 제한
+            low_cpu_mem_usage=True
         )
         print("   ✅ 커스텀 모델 로딩 완료")
     else:
@@ -116,7 +115,8 @@ def chat_with_model(message, history, image=None):
                 pixel_values = transform(pil_image).unsqueeze(0)
                 image_metas = {"vision_grid_thw": torch.tensor([[1, 14, 14]])}  # 기본 그리드 크기
                 
-                outputs = model.generate(
+                # 멀티모달 모델의 forward 메서드 사용
+                outputs = model(
                     input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"],
                     pixel_values=[pixel_values],
@@ -128,7 +128,7 @@ def chat_with_model(message, history, image=None):
                 )
             else:
                 # 이미지가 없는 경우 텍스트만 생성
-                outputs = model.generate(
+                outputs = model(
                     input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"],
                     max_new_tokens=200,
@@ -137,7 +137,38 @@ def chat_with_model(message, history, image=None):
                     pad_token_id=tokenizer.eos_token_id
                 )
         
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # outputs가 튜플인 경우 첫 번째 요소 사용
+        if isinstance(outputs, tuple):
+            logits = outputs[0]
+        else:
+            logits = outputs.logits if hasattr(outputs, 'logits') else outputs
+        
+        # 가장 높은 확률의 토큰 선택
+        next_token = torch.argmax(logits[:, -1, :], dim=-1)
+        generated_tokens = [next_token]
+        
+        # 추가 토큰 생성
+        for _ in range(199):  # max_new_tokens - 1
+            inputs["input_ids"] = torch.cat([inputs["input_ids"], next_token.unsqueeze(-1)], dim=-1)
+            inputs["attention_mask"] = torch.cat([inputs["attention_mask"], torch.ones_like(next_token.unsqueeze(-1))], dim=-1)
+            
+            with torch.no_grad():
+                outputs = model(**inputs)
+                if isinstance(outputs, tuple):
+                    logits = outputs[0]
+                else:
+                    logits = outputs.logits if hasattr(outputs, 'logits') else outputs
+                
+                next_token = torch.argmax(logits[:, -1, :], dim=-1)
+                generated_tokens.append(next_token)
+                
+                if next_token.item() == tokenizer.eos_token_id:
+                    break
+        
+        # 생성된 토큰들을 디코딩
+        generated_ids = torch.cat(generated_tokens, dim=0)
+        response = tokenizer.decode(generated_ids, skip_special_tokens=True)
+        
         if message in response:
             response = response.replace(message, "").strip()
         return response if response else "죄송합니다. 응답을 생성할 수 없습니다."
@@ -172,7 +203,8 @@ def solve_math_problem(problem, image=None):
                 pixel_values = transform(pil_image).unsqueeze(0)
                 image_metas = {"vision_grid_thw": torch.tensor([[1, 14, 14]])}  # 기본 그리드 크기
                 
-                outputs = model.generate(
+                # 멀티모달 모델의 forward 메서드 사용
+                outputs = model(
                     input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"],
                     pixel_values=[pixel_values],
@@ -184,7 +216,7 @@ def solve_math_problem(problem, image=None):
                 )
             else:
                 # 이미지가 없는 경우 텍스트만 생성
-                outputs = model.generate(
+                outputs = model(
                     input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"],
                     max_new_tokens=300,
@@ -193,7 +225,38 @@ def solve_math_problem(problem, image=None):
                     pad_token_id=tokenizer.eos_token_id
                 )
         
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # outputs가 튜플인 경우 첫 번째 요소 사용
+        if isinstance(outputs, tuple):
+            logits = outputs[0]
+        else:
+            logits = outputs.logits if hasattr(outputs, 'logits') else outputs
+        
+        # 가장 높은 확률의 토큰 선택
+        next_token = torch.argmax(logits[:, -1, :], dim=-1)
+        generated_tokens = [next_token]
+        
+        # 추가 토큰 생성
+        for _ in range(299):  # max_new_tokens - 1
+            inputs["input_ids"] = torch.cat([inputs["input_ids"], next_token.unsqueeze(-1)], dim=-1)
+            inputs["attention_mask"] = torch.cat([inputs["attention_mask"], torch.ones_like(next_token.unsqueeze(-1))], dim=-1)
+            
+            with torch.no_grad():
+                outputs = model(**inputs)
+                if isinstance(outputs, tuple):
+                    logits = outputs[0]
+                else:
+                    logits = outputs.logits if hasattr(outputs, 'logits') else outputs
+                
+                next_token = torch.argmax(logits[:, -1, :], dim=-1)
+                generated_tokens.append(next_token)
+                
+                if next_token.item() == tokenizer.eos_token_id:
+                    break
+        
+        # 생성된 토큰들을 디코딩
+        generated_ids = torch.cat(generated_tokens, dim=0)
+        response = tokenizer.decode(generated_ids, skip_special_tokens=True)
+        
         if prompt in response:
             response = response.replace(prompt, "").strip()
         return response if response else "죄송합니다. 수학 문제를 풀 수 없습니다."