.

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import spaces
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, AutoTokenizer
 import torch
 import logging
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # Configure logging/logger
 logging.basicConfig(
@@ -18,6 +19,7 @@ pipelines = {}
 # Predefined list of models to compare (can be expanded)
 model_options = {
     "Foundation-Sec-8B": "fdtn-ai/Foundation-Sec-8B",
+    "Llama-3.1-8B": "meta-llama/Llama-3.1-8B",
 }
 
 # Initialize models at startup
@@ -75,11 +77,10 @@ def create_demo():
 
         # Input Section
         with gr.Row():
-            system_message = gr.Textbox(
+            prompt = gr.Textbox(
                 value="You are a helpful assistant providing answers for technical and customer support queries.",
-                label="System message"
+                label="Prompt"
             )
-            user_message = gr.Textbox(label="Your question", placeholder="Type your question here...")
 
         with gr.Row():
             max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
@@ -97,38 +98,51 @@ def create_demo():
 
         # Dynamic Response Section
         response_box1 = gr.Textbox(label="Response from Model 1", interactive=False)
-        #response_box2 = gr.Textbox(label="Response from Model 2", interactive=False)
+        response_box2 = gr.Textbox(label="Response from Model 2", interactive=False)
 
         # Function to generate responses
         def generate_responses(
-            message, system_message, max_tokens, temperature, top_p, selected_models
+            message, prompt, max_tokens, temperature, top_p, selected_models
         ):
-            #if len(selected_models) != 2:
-            #    return "Error: Please select exactly two models to compare.", ""
+            if len(selected_models) != 2:
+                return "Error: Please select exactly two models to compare.", ""
             
             if len(selected_models) == 0:
-                return "Error: Please select at least one model"
+                return "Error: Please select at least one model", ""
             
-            model_path = model_options[selected_models[0]]
-            full_prompt = f"{system_message}\n\nUser: {message}\nAssistant:"
-            response = generate_text_local(
-                model_path,
-                full_prompt,
-                max_tokens,
-                temperature,
-                top_p
-            )
+            # 選択されたモデルの結果を格納する辞書
+            responses = {}
+            futures_to_model = {}  # 各futureとモデルを紐づけるための辞書
+                        
+            with ThreadPoolExecutor(max_workers=len(selected_models)) as executor:
+                # 各モデルに対してタスクを提出
+                futures = []
+                for model_name in selected_models:
+                    model_path = model_options[model_name]
+                    future = executor.submit(
+                        generate_text_local, 
+                        model_path, 
+                        prompt, 
+                        max_tokens, 
+                        temperature, 
+                        top_p
+                    )
+                    futures.append(future)
+                    futures_to_model[future] = model_name
+                
+                # 結果の収集
+                for future in as_completed(futures):
+                    model_name = futures_to_model[future]
+                    responses[model_name] = future.result()
             
-            #return responses.get(selected_models[0], ""), responses.get(selected_models[1], "")
-            return response
+            return responses.get(selected_models[0], ""), responses.get(selected_models[1], "")
             
         # Add a button for generating responses
         submit_button = gr.Button("Generate Responses")
         submit_button.click(
             generate_responses,
-            inputs=[user_message, system_message, max_tokens, temperature, top_p, selected_models],
-            #outputs=[response_box1, response_box2],  # Link to response boxes
-            outputs=[response_box1]
+            inputs=[prompt, max_tokens, temperature, top_p, selected_models],
+            outputs=[response_box1, response_box2],  # Link to response boxes
         )
 
     return demo