Update app.py

app.py

@@ -1,60 +1,62 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import time
-import random
+import pandas as pd
+from utils.dataset_loader import load_dataset
 
-# Load the models and tokenizer
-model_name = "Canstralian/text2shellcommands"  # Choose your model, can be changed based on use case
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load dataset
+df = load_dataset()
 
-# Function to generate shell command or response based on the prompt
-def generate_shell_command(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=50, num_return_sequences=1)
-    command = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return command
+# Function to simulate conversation with model selection
+def chat_interface(user_input, selected_model, prompt_id=None):
+    if df is not None and prompt_id is not None:
+        prompt = df.iloc[prompt_id]["prompt_text"]  # Replace with the actual column name
+        response = f"[{selected_model}] used the debugging prompt: '{prompt}'.\nUser said: '{user_input}'\nResponse: Simulated output."
+    else:
+        response = f"[{selected_model}] says: You entered '{user_input}'. This is a simulated response."
+    return response
 
-# Function to simulate a retro terminal environment with some 90s hacker vibe
-def terminal_ui(prompt):
-    # Simulate typing effect
-    fake_typing_effect = [
-        "Initializing...\n",
-        "Boot sequence complete...\n",
-        "Connecting to secure network...\n",
-        "Accessing restricted files...\n",
-        "Running diagnostics...\n",
-        "Command input: " + prompt + "\n"
-    ]
-    
-    # Adding some suspense and random time delays to create that 'hacker' feel
-    for line in fake_typing_effect:
-        time.sleep(random.uniform(0.5, 1.5))  # Simulate typing delay
-        print(line)  # Simulate print to terminal
-        time.sleep(0.3)
-    
-    # Get AI-generated response for the command prompt
-    command_response = generate_shell_command(prompt)
-    
-    # Simulate result display with some retro terminal feedback
-    result_output = f"\n[ SYSTEM STATUS: OK ]\n[ {random.choice(['OK', 'ERROR', 'WARNING'])} ]\n\n"
-    result_output += f"Command executed: {command_response}\n"
-    result_output += "[ End of output ]"
-    
-    return result_output
+# List of available models (Updated as per your request)
+models = ["Canstralian/text2shellcommands", "Canstralian/RabbitRedux", "Canstralian/CySec_Known_Exploit_Analyzer"]
+prompt_ids = df.index.tolist() if df is not None else []
 
-# Create a Gradio interface with a retro terminal design
-def retro_terminal_interface(prompt):
-    result = terminal_ui(prompt)
-    return result
+# Gradio Interface
+with gr.Blocks(css="./static/styles.css") as demo:
+    with gr.Row():
+        gr.Markdown("### Retro Hacker Chat with Debugging Prompts", elem_classes="retro-terminal")
+    with gr.Row():
+        user_input = gr.Textbox(
+            label="Enter your message:",
+            placeholder="Type your message here...",
+            elem_classes="retro-terminal"
+        )
+        model_selector = gr.Dropdown(
+            choices=models,
+            label="Select Model",
+            value=models[0],
+            elem_classes="retro-terminal"
+        )
+        if prompt_ids:
+            prompt_selector = gr.Dropdown(
+                choices=prompt_ids,
+                label="Select Debugging Prompt ID",
+                value=prompt_ids[0],
+                elem_classes="retro-terminal"
+            )
+        else:
+            prompt_selector = None
+    with gr.Row():
+        response_box = gr.Textbox(
+            label="Model Response:",
+            placeholder="The model's response will appear here...",
+            elem_classes="retro-terminal"
+        )
+    with gr.Row():
+        send_button = gr.Button("Send", elem_classes="retro-terminal")
 
-# Launch the Gradio app with a terminal theme
-iface = gr.Interface(
-    fn=retro_terminal_interface, 
-    inputs=gr.Textbox(placeholder="Type your shell command here...", label="Enter Command:"),
-    outputs=gr.Textbox(label="Terminal Output", lines=20, interactive=False),
-    theme="compact",  # Use Gradio's built-in compact theme for a terminal-like feel
-    live=True  # Enable live feedback to simulate a real-time terminal experience
-)
+    # Link input and output
+    if prompt_selector:
+        send_button.click(chat_interface, inputs=[user_input, model_selector, prompt_selector], outputs=response_box)
+    else:
+        send_button.click(chat_interface, inputs=[user_input, model_selector], outputs=response_box)
 
-iface.launch()
+# Launch the interface
+demo.launch()