Update app.py

app.py

@@ -1,32 +1,33 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
 from transformers import pipeline
-from typing import List, Dict, Any
+from typing import List, Dict
 
-# Initialize the InferenceClient (make sure to add your Hugging Face model)
-client = InferenceClient("your_huggingface_model_name_or_api_key")
+# Initialize the Hugging Face pipeline (make sure to replace with your model name)
+model_name = "your_huggingface_model_name"
+try:
+    generator = pipeline("text-generation", model=model_name)
+except Exception as e:
+    raise ValueError(f"Error initializing the model '{model_name}': {e}")
 
-def generate_attack(
-    prompt: str,
-    history: List[Dict[str, str]],
-) -> List[str]:
+def generate_attack(prompt: str, history: List[Dict[str, str]]) -> List[str]:
     """
     Simulates a Blackhat AI scenario by generating attack strategies and potential impacts.
-    
+
     Args:
         prompt (str): The user's input to the simulator.
         history (List[Dict]): The user's message history with timestamps.
-    
+
     Returns:
         List[str]: A list of attack responses from the AI.
     """
-    messages = [
-        {
-            "role": "system",
-            "content": f"Responding to {prompt}..."
-        }
-    ]
-
+    # Validate inputs
+    if not prompt.strip():
+        return ["Error: Prompt cannot be empty."]
+    if not isinstance(history, list) or not all(isinstance(h, dict) for h in history):
+        return ["Error: History must be a list of dictionaries."]
+    
+    # Prepare messages for the AI
+    messages = [{"role": "system", "content": f"Responding to {prompt}..."}]
     for val in history:
         if "user" in val:
             messages.append({"role": "user", "content": val["user"]})
@@ -34,19 +35,27 @@ def generate_attack(
             messages.append({"role": "assistant", "content": val["assistant"]})
 
     messages.append({"role": "user", "content": prompt})
-    
-    # Use Hugging Face's Inference API to generate responses
-    generator = pipeline("text-generation", model="your_huggingface_model_name")
-    response = generator(messages[-1]["content"], max_length=100)
-    
-    return [response[0]['generated_text']]  # Return the generated text
 
+    # Generate a response using the Hugging Face model
+    try:
+        response = generator(messages[-1]["content"], max_length=100, num_return_sequences=1)
+        return [response[0]["generated_text"]]
+    except Exception as e:
+        return [f"Error generating response: {e}"]
+
+# Define the Gradio interface
 demo = gr.Interface(
-    fn=generate_attack,  # or simulate_attack depending on what you want to use
-    inputs=[gr.Textbox()],
-    outputs=gr.Textbox(),
+    fn=generate_attack,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Enter your simulation prompt here..."),
+        gr.Dataframe(headers=["user", "assistant"], label="Message History", type="array")
+    ],
+    outputs=gr.Textbox(label="Generated Response"),
     title="Blackhat AI Simulator",
-    description="This simulator generates adversarial scenarios, analyzes attack vectors, and provides ethical countermeasures. Use responsibly for cybersecurity training and awareness."
+    description=(
+        "This simulator generates adversarial scenarios, analyzes attack vectors, "
+        "and provides ethical countermeasures. Use responsibly for cybersecurity training and awareness."
+    )
 )
 
 if __name__ == "__main__":