Update app.py

app.py

@@ -3,6 +3,7 @@ from llama_cpp import Llama
 
 st.set_page_config(page_title="Cybertron Chat Interface", layout="wide")
 st.title("🧠 Cybertron Chat: Generalist vs Specialist Mode")
+
 st.markdown("""
 Welcome to the Cybertron Chat Interface. Choose between:
 
@@ -10,10 +11,23 @@ Welcome to the Cybertron Chat Interface. Choose between:
 - 🛡️ **Specialist Mode** for precise, tactical cyber security and pentesting insights
 """)
 
-# Sidebar toggle for mode
-model_choice = st.sidebar.radio("Choose Mode:", ["Generalist 🤖", "Specialist 🛡️"], help="Switch between general reasoning and focused cybersec models")
+# -- Sidebar toggle for mode --
+model_choice = st.sidebar.radio(
+    "Choose Mode:", 
+    ["Generalist 🤖", "Specialist 🛡️"], 
+    help="Switch between general reasoning and focused cybersec models"
+)
+
+# -- Load model only once with caching --
+@st.cache_resource(show_spinner=False)
+def load_model(repo_id, filename):
+    return Llama.from_pretrained(
+        repo_id=repo_id,
+        filename=filename,
+        n_ctx=2048,  # Default context window
+    )
 
-# Display model tips
+# -- Set mode-specific models and prompts --
 if model_choice == "Generalist 🤖":
     st.sidebar.markdown("""
     **Best for:**
@@ -21,15 +35,16 @@ if model_choice == "Generalist 🤖":
     - Brainstorming tools or strategies
     - Simulating conversations
     """)
-    llm = Llama.from_pretrained(
+    llm = load_model(
         repo_id="bartowski/cybertron-v4-qw7B-MGS-GGUF",
-        filename="cybertron-v4-qw7B-MGS-IQ2_M.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"
     ]
+
 else:
     st.sidebar.markdown("""
     **Best for:**
@@ -37,9 +52,9 @@ else:
     - Red/blue team planning
     - Shell scripting and command-line tasks
     """)
-    llm = Llama.from_pretrained(
+    llm = load_model(
         repo_id="TheBloke/una-cybertron-7B-v2-GGUF",
-        filename="una-cybertron-7b-v2-bf16.Q2_K.gguf",
+        filename="una-cybertron-7b-v2-bf16.Q2_K.gguf"
     )
     example_prompts = [
         "List enumeration commands for Active Directory",
@@ -47,21 +62,27 @@ else:
         "Generate a Bash reverse shell with obfuscation"
     ]
 
-# Prompt input and example selector
+# -- Prompt input and example selector --
 col1, col2 = st.columns([2, 1])
 
 with col1:
-    user_input = st.text_area("\U0001F4AC Enter your query below:", height=150)
+    user_input = st.text_area("💬 Enter your query below:", height=150)
 with col2:
-    st.markdown("**\U0001F4D6 Prompt Examples:**")
+    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
 
-# Run inference
+# -- Run inference --
 if st.button("Submit", use_container_width=True):
-    with st.spinner("Generating response..."):
-        output = llm(user_input, max_tokens=512, echo=True)
-        st.markdown("---")
-        st.markdown("**\U0001F4C4 Response:**")
-        st.code(output["choices"][0]["text"].strip())
+    if user_input.strip() == "":
+        st.warning("Please enter a prompt or select an example.")
+    else:
+        with st.spinner("🚀 Generating response..."):
+            try:
+                output = llm(user_input, max_tokens=512, echo=False)
+                st.markdown("---")
+                st.markdown("**📄 Response:**")
+                st.code(output["choices"][0]["text"].strip())
+            except Exception as e:
+                st.error(f"⚠️ Error: {str(e)}")