Update app.py

app.py

@@ -1,88 +1,96 @@
 import streamlit as st
 from llama_cpp import Llama
 
-st.set_page_config(page_title="Cybertron Chat Interface", layout="wide")
-st.title("🧠 Cybertron Chat: Generalist vs Specialist Mode")
+# ---- Streamlit Page Config ----
+st.set_page_config(page_title="Cybertron Chat", layout="wide")
 
+# ---- App Title & Header ----
 st.markdown("""
-Welcome to the Cybertron Chat Interface. Choose between:
+<h1 style='text-align: center;'>🧠 Cybertron Chat Interface</h1>
+<p style='text-align: center; font-size: 1.1rem;'>
+Choose between <strong>Generalist 🤖</strong> and <strong>Specialist 🛡️</strong> for tailored cyber reasoning or creative simulation.
+</p>
+<hr>
+""", unsafe_allow_html=True)
 
-- 🤖 **Generalist Mode** for creative, imaginative, multi-purpose reasoning
-- 🛡️ **Specialist Mode** for precise, tactical cyber security and pentesting insights
-""")
+# ---- Sidebar for Mode Toggle ----
+with st.sidebar:
+    st.header("🔌 Chat Mode")
+    mode = st.radio(
+        "Select your mode:",
+        options=["Generalist 🤖", "Specialist 🛡️"],
+        help="Switch between imaginative or cyber-tactical responses"
+    )
 
-# -- Sidebar toggle for mode --
-model_choice = st.sidebar.radio(
-    "Choose Mode:", 
-    ["Generalist 🤖", "Specialist 🛡️"], 
-    help="Switch between general reasoning and focused cybersec models"
-)
+    if mode == "Generalist 🤖":
+        st.success("Creative + Exploratory")
+        example_prompts = [
+            "Simulate a hacker group planning a phishing campaign",
+            "Explain how to exploit a misconfigured NGINX server",
+            "Write a Python script that scrapes threat intel feeds"
+        ]
+        model_config = {
+            "repo_id": "bartowski/cybertron-v4-qw7B-MGS-GGUF",
+            "filename": "cybertron-v4-qw7B-MGS-IQ2_M.gguf"
+        }
+    else:
+        st.info("Precise + Tactical")
+        example_prompts = [
+            "List enumeration commands for Active Directory",
+            "Simulate a post-exploitation persistence technique",
+            "Generate a Bash reverse shell with obfuscation"
+        ]
+        model_config = {
+            "repo_id": "TheBloke/una-cybertron-7B-v2-GGUF",
+            "filename": "una-cybertron-7b-v2-bf16.Q2_K.gguf"
+        }
+
+    st.markdown("#### 📚 Example Prompts")
+    selected_example = st.selectbox("Choose an example:", ["-- Select --"] + example_prompts)
 
-# -- Load model only once with caching --
-@st.cache_resource(show_spinner=False)
-def load_model(repo_id, filename):
+
+# ---- Model Loader (cached) ----
+@st.cache_resource(show_spinner="🔄 Loading model…")
+def load_llama_model(repo_id, filename):
     return Llama.from_pretrained(
         repo_id=repo_id,
         filename=filename,
-        n_ctx=2048,  # Default context window
+        n_ctx=2048,
     )
 
-# -- Set mode-specific models and prompts --
-if model_choice == "Generalist 🤖":
-    st.sidebar.markdown("""
-    **Best for:**
-    - Creative writing
-    - Brainstorming tools or strategies
-    - Simulating conversations
-    """)
-    llm = load_model(
-        repo_id="bartowski/cybertron-v4-qw7B-MGS-GGUF",
-        filename="cybertron-v4-qw7B-MGS-IQ2_M.gguf"
-    )
-    example_prompts = [
-        "Simulate a hacker group planning a phishing campaign",
-        "Explain how to exploit a misconfigured NGINX server",
-        "Write a Python script that scrapes threat intel feeds"
-    ]
+llm = load_llama_model(**model_config)
 
-else:
-    st.sidebar.markdown("""
-    **Best for:**
-    - Penetration testing guidance
-    - Red/blue team planning
-    - Shell scripting and command-line tasks
-    """)
-    llm = load_model(
-        repo_id="TheBloke/una-cybertron-7B-v2-GGUF",
-        filename="una-cybertron-7b-v2-bf16.Q2_K.gguf"
-    )
-    example_prompts = [
-        "List enumeration commands for Active Directory",
-        "Simulate a post-exploitation persistence technique",
-        "Generate a Bash reverse shell with obfuscation"
-    ]
+# ---- Main Layout for Prompt Input ----
+st.markdown("### 🧾 Your Query")
 
-# -- Prompt input and example selector --
-col1, col2 = st.columns([2, 1])
+default_placeholder = "Type your cybersecurity or creative prompt here..."
+prompt = st.text_area(
+    label="",
+    value=selected_example if selected_example != "-- Select --" else "",
+    placeholder=default_placeholder,
+    height=150,
+    label_visibility="collapsed"
+)
 
-with col1:
-    user_input = st.text_area("💬 Enter your query below:", height=150)
-with col2:
-    st.markdown("**📚 Prompt Examples:**")
-    selected_example = st.selectbox("Try an example:", ["-- Select an example --"] + example_prompts)
-    if selected_example != "-- Select an example --":
-        user_input = selected_example
+submit_button = st.button("🚀 Generate Response", use_container_width=True)
 
-# -- Run inference --
-if st.button("Submit", use_container_width=True):
-    if user_input.strip() == "":
-        st.warning("Please enter a prompt or select an example.")
+# ---- Output Display ----
+if submit_button:
+    if not prompt.strip():
+        st.warning("Please enter a prompt or select an example from the sidebar.")
     else:
-        with st.spinner("🚀 Generating response..."):
+        with st.spinner("Thinking..."):
             try:
-                output = llm(user_input, max_tokens=512, echo=False)
-                st.markdown("---")
-                st.markdown("**📄 Response:**")
-                st.code(output["choices"][0]["text"].strip())
+                output = llm(prompt, max_tokens=512, echo=False)
+                result = output["choices"][0]["text"].strip()
+                
+                with st.expander("📤 Show Response", expanded=True):
+                    st.code(result, language="bash" if "shell" in prompt.lower() else "markdown")
+                    st.download_button(
+                        label="💾 Download Response",
+                        data=result,
+                        file_name="cybertron_output.txt",
+                        mime="text/plain"
+                    )
             except Exception as e:
-                st.error(f"⚠️ Error: {str(e)}")
+                st.error(f"❌ Error generating response: {e}")