.

app.py

@@ -6,6 +6,7 @@ from enum import Enum, auto
 import torch
 from transformers import AutoTokenizer, pipeline
 import spaces
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # ロガーの設定
 logging.basicConfig(
@@ -114,12 +115,15 @@ def generate_text_local(model_path, text):
     """ローカルモデルでのテキスト生成"""
     try:
         logger.info(f"Running local text generation with {model_path}")
-        outputs = pipelines[model_path](
+        pipeline = pipelines[model_path]
+        pipeline.to("cuda")
+        outputs = pipeline(
             text,
             max_new_tokens=40,
             do_sample=False,
             num_return_sequences=1
         )
+        pipeline.to("cpu")
         return outputs[0]["generated_text"]
     except Exception as e:
         logger.error(f"Error in local text generation with {model_path}: {str(e)}")
@@ -144,7 +148,10 @@ def classify_text_local(model_path, text):
     """ローカルモデルでのテキスト分類"""
     try:
         logger.info(f"Running local classification with {model_path}")
-        result = pipelines[model_path](text)
+        pipeline = pipelines[model_path]
+        pipeline.to("cuda")
+        result = pipeline(text)
+        pipeline.to("cpu")
         return str(result)
     except Exception as e:
         logger.error(f"Error in local classification with {model_path}: {str(e)}")
@@ -164,28 +171,43 @@ def classify_text_api(model_id, text):
 def handle_invoke(text, selected_types):
     """Invokeボタンのハンドラ"""
     results = []
+    futures_to_model = {}  # 各futureとモデルを紐づけるための辞書
     
-    # テキスト生成モデルの実行
-    for model in TEXT_GENERATION_MODELS:
-        if model["type"] in selected_types:
-            if model["type"] == LOCAL:
-                result = generate_text_local(model["model_path"], text)
-            else:  # api
-                result = generate_text_api(model["model_id"], text)
-            results.append(f"{model['name']}: {result}")
-    
-    # 分類モデルの実行
-    for model in CLASSIFICATION_MODELS:
-        if model["type"] in selected_types:
-            if model["type"] == LOCAL:
-                result = classify_text_local(model["model_path"], text)
-            else:  # api
-                result = classify_text_api(model["model_id"], text)
-            results.append(f"{model['name']}: {result}")
+    with ThreadPoolExecutor(max_workers=len(selected_types)) as executor:
+        futures = []
+        
+        # テキスト生成モデルの実行
+        for model in TEXT_GENERATION_MODELS:
+            if model["type"] in selected_types:
+                if model["type"] == LOCAL:
+                    future = executor.submit(generate_text_local, model["model_path"], text)
+                    futures.append(future)
+                    futures_to_model[future] = model
+                else:  # api
+                    future = executor.submit(generate_text_api, model["model_id"], text)
+                    futures.append(future)
+                    futures_to_model[future] = model
     
-    # 結果リストの長さを調整
-    while len(results) < len(TEXT_GENERATION_MODELS) + len(CLASSIFICATION_MODELS):
-        results.append("")
+        # 分類モデルの実行
+        for model in CLASSIFICATION_MODELS:
+            if model["type"] in selected_types:
+                if model["type"] == LOCAL:
+                    future = executor.submit(classify_text_local, model["model_path"], text)
+                    futures.append(future)
+                    futures_to_model[future] = model
+                else:  # api
+                    future = executor.submit(classify_text_api, model["model_id"], text)
+                    futures.append(future)
+                    futures_to_model[future] = model
+        
+        # 結果の収集（モデルの順序を保持）
+        all_models = TEXT_GENERATION_MODELS + CLASSIFICATION_MODELS
+        results = [""] * len(all_models)  # 事前に結果リストを初期化
+        
+        for future in as_completed(futures):
+            model = futures_to_model[future]
+            model_index = all_models.index(model)
+            results[model_index] = future.result()
     
     return results