Update app.py

app.py

@@ -0,0 +1,38 @@
+import gradio as gr
+from transformers import pipeline
+
+# Load model (change to your own if needed)
+llm = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", max_new_tokens=300)
+
+# Role-based context
+roles = {
+    "Blue Team": "You are a cybersecurity analyst defending a system.",
+    "Red Team": "You are simulating a cyber attack for penetration testing.",
+    "GRC": "You are an AI compliance auditor evaluating AI systems."
+}
+
+# Generation logic
+def generate_response(prompt, role):
+    context = roles.get(role, "")
+    full_prompt = f"{context}\n{prompt}"
+    try:
+        result = llm(full_prompt)[0]['generated_text']
+    except Exception as e:
+        result = f"Error: {str(e)}"
+    return result
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🛡️ CyberAI Prompt Generator")
+    gr.Markdown("Select your role and input a prompt to get a contextual AI response.")
+
+    with gr.Row():
+        role = gr.Dropdown(choices=list(roles.keys()), label="Role", value="Blue Team")
+        prompt = gr.Textbox(label="Your prompt", placeholder="e.g. Generate a threat model for a hospital IoT network")
+
+    output = gr.Textbox(label="AI Response")
+
+    btn = gr.Button("Generate")
+    btn.click(fn=generate_response, inputs=[prompt, role], outputs=output)
+
+demo.launch()