.

app.py

@@ -117,9 +117,13 @@ def generate_text_local(model_path, text):
         logger.info(f"Running local text generation with {model_path}")
         pipeline = pipelines[model_path]
         
-        # モデルを明示的に GPU に移動
-        if hasattr(pipeline.model, "to"):
-            pipeline.model.to("cuda")
+        # モデルをGPUに移動（パイプライン全体）
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        pipeline.model = pipeline.model.to(device)
+        
+        # トークナイザーなど他のパイプラインコンポーネントがGPUを使用するように設定
+        if hasattr(pipeline, "device"):
+            pipeline.device = device
         
         # デバイス情報をログに記録
         device_info = next(pipeline.model.parameters()).device
@@ -129,13 +133,15 @@ def generate_text_local(model_path, text):
             text,
             max_new_tokens=40,
             do_sample=False,
-            num_return_sequences=1
+            num_return_sequences=1,
+            device=device  # デバイスを明示的に指定
         )
-
-        # モデルを CPU に戻す
-        if hasattr(pipeline.model, "to"):
-            pipeline.model.to("cpu")
-
+        
+        # モデルをCPUに戻す
+        pipeline.model = pipeline.model.to("cpu")
+        if hasattr(pipeline, "device"):
+            pipeline.device = torch.device("cpu")
+        
         return outputs[0]["generated_text"]
     except Exception as e:
         logger.error(f"Error in local text generation with {model_path}: {str(e)}")
@@ -162,20 +168,25 @@ def classify_text_local(model_path, text):
         logger.info(f"Running local classification with {model_path}")
         pipeline = pipelines[model_path]
         
-        # モデルを明示的に GPU に移動
-        if hasattr(pipeline.model, "to"):
-            pipeline.model.to("cuda")
-
+        # モデルをGPUに移動（パイプライン全体）
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        pipeline.model = pipeline.model.to(device)
+        
+        # トークナイザーなど他のパイプラインコンポーネントがGPUを使用するように設定
+        if hasattr(pipeline, "device"):
+            pipeline.device = device
+        
         # デバイス情報をログに記録
         device_info = next(pipeline.model.parameters()).device
         logger.info(f"Model {model_path} is running on device: {device_info}")
         
-        result = pipeline(text)
-
-        # モデルを CPU に戻す
-        if hasattr(pipeline.model, "to"):
-            pipeline.model.to("cpu")
-
+        result = pipeline(text, device=device)  # デバイスを明示的に指定
+        
+        # モデルをCPUに戻す
+        pipeline.model = pipeline.model.to("cpu")
+        if hasattr(pipeline, "device"):
+            pipeline.device = torch.device("cpu")
+            
         return str(result)
     except Exception as e:
         logger.error(f"Error in local classification with {model_path}: {str(e)}")