.

README.md

@@ -1,6 +1,6 @@
 ---
-title: Compare Models
-emoji: 🐨
+title: Compare Security Models
+emoji: 🐼
 colorFrom: gray
 colorTo: yellow
 sdk: gradio

app.py

@@ -1,52 +1,58 @@
-import os
 import gradio as gr
-from huggingface_hub import login, InferenceClient
 import spaces
+from transformers import pipeline
+import torch
+import logging
 
-# Authenticate with Hugging Face API
-api_key = os.getenv("TOKEN")
-login(api_key)
+# Configure logging/logger
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
 
 # Predefined list of models to compare (can be expanded)
 model_options = {
-    "Llama-3.1-70B": "meta-llama/Llama-3.1-70B-Instruct",
-    "Qwen-2.5-1.5B-Instruct": "Qwen/Qwen2.5-1.5B-Instruct",
-    "Llama-3.2-1B": "meta-llama/Llama-3.2-1B",
-    "DeepSeek-V2.5": "deepseek-ai/DeepSeek-V2.5",
-    "Athene-V2-Chat": "Nexusflow/Athene-V2-Chat",
+    "Foundation-Sec-8B": pipeline("text-generation", model="fdtn-ai/Foundation-Sec-8B"),
 }
 
-# Initialize clients for models
-clients = {name: InferenceClient(repo_id) for name, repo_id in model_options.items()}
-
 # Define the response function
 @spaces.GPU
-def respond(
-    message,
-    history: list[dict],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    selected_models,
-):
-    messages = [{"role": "system", "content": system_message}] + history
-    messages.append({"role": "user", "content": message})
-
-    responses = {}
-
-    # Generate responses for each selected model
-    for model_name in selected_models:
-        client = clients[model_name]
-        response = ""
-        for token in client.chat_completion(
-            messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
-        ):
-            delta = token.choices[0].delta.content
-            response += delta
-        responses[model_name] = response
+def generate_text_local(model_pipeline, prompt):
+    """Local text generation"""
+    try:
+        logger.info(f"Running local text generation with {model_pipeline.path}")
+        
+        # Move model to GPU (entire pipeline)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model_pipeline.model = model_pipeline.model.to(device)
+        
+        # Set other pipeline components to use GPU
+        if hasattr(model_pipeline, "device"):
+            model_pipeline.device = device
+        
+        # Record device information
+        device_info = next(model_pipeline.model.parameters()).device
+        logger.info(f"Model {model_pipeline.path} is running on device: {device_info}")
+        
+        outputs = model_pipeline(
+            prompt,
+            max_new_tokens=3,      # = model.generate(max_new_tokens=3, …)
+            do_sample=True,
+            temperature=0.1,
+            top_p=0.9,
+            clean_up_tokenization_spaces=True,  # echo 部分を整形
+        )
 
-    return responses
+        # Move model back to CPU
+        model_pipeline.model = model_pipeline.model.to("cpu")
+        if hasattr(model_pipeline, "device"):
+            model_pipeline.device = torch.device("cpu")
+        
+        return outputs[0]["generated_text"].replace(prompt, "").strip()
+    except Exception as e:
+        logger.error(f"Error in local text generation with {model_pipeline.path}: {str(e)}")
+        return f"Error: {str(e)}"
 
 # Build Gradio app
 def create_demo():
@@ -77,31 +83,34 @@ def create_demo():
         # Model Selection Section
         selected_models = gr.CheckboxGroup(
             choices=list(model_options.keys()),
-            label="Select exactly two models to compare",
-            value=["Llama-3.1-70B", "Qwen-2.5-1.5B-Instruct"],  # Default models
+            label="Select exactly two model to compare",
+            value=["Foundation-Sec-8B"],  # Default models
         )
 
         # 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
         ):
-            if len(selected_models) != 2:
-                return "Error: Please select exactly two models to compare.", ""
-            responses = respond(
-                message, [], system_message, max_tokens, temperature, top_p, selected_models
+            #if len(selected_models) != 2:
+            #    return "Error: Please select exactly two models to compare.", ""
+            responses = generate_text_local(
+                #message, [], system_message, max_tokens, temperature, top_p, selected_models
+                selected_models[0],
+                message
             )
-            return responses.get(selected_models[0], ""), responses.get(selected_models[1], "")
-
+            #return responses.get(selected_models[0], ""), responses.get(selected_models[1], "")
+            return responses
         # 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, response_box2],  # Link to response boxes
+            outputs=[response_box1]
         )
 
     return demo