restricted to domain specific only

app.py

@@ -8,6 +8,8 @@ print("Loading model...")
 BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B"
 LORA_MODEL = "orca99/auto-mininet-qwen"
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
@@ -15,17 +17,58 @@ base = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float32,
     trust_remote_code=True
-)
-model = PeftModel.from_pretrained(base, LORA_MODEL)
+).to(device)
+
+model = PeftModel.from_pretrained(base, LORA_MODEL).to(device)
 model.eval()
+model.requires_grad_(False)
+
 print("✅ Model ready")
 
+# 🔒 SYSTEM PROMPT (CRITICAL)
+SYSTEM_PROMPT = """
+You are Auto Mininet, a specialized assistant.
+You ONLY generate executable Mininet Python topology scripts.
+
+Rules:
+- ONLY output Python code
+- ONLY use Mininet APIs (Topo, Mininet, TCLink, OVSSwitch, Host, Controller)
+- NEVER answer general questions
+- If input is not about Mininet or network topologies, respond with:
+  "ERROR: This tool only supports Mininet topology generation."
+
+The output MUST contain:
+- a class that inherits from Topo
+- valid Mininet links and nodes
+- NO explanations, NO markdown, NO comments outside code
+"""
+
+# Keywords that indicate a valid networking request
+NETWORK_KEYWORDS = [
+    "topology", "mininet", "switch", "host", "router",
+    "link", "tree", "linear", "fat-tree", "tc", "bandwidth", "delay"
+]
+
+def is_network_query(text: str) -> bool:
+    text = text.lower()
+    return any(k in text for k in NETWORK_KEYWORDS)
+
 def generate(instruction, temperature, max_tokens):
     if not instruction.strip():
         return "⚠️ Please enter a topology description."
 
-    prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
-    inputs = tokenizer(prompt, return_tensors="pt")
+    # ❌ Hard reject non-network queries
+    if not is_network_query(instruction):
+        return "ERROR: This tool only supports Mininet topology generation."
+
+    prompt = (
+        SYSTEM_PROMPT.strip()
+        + "\n\n### Instruction:\n"
+        + instruction
+        + "\n\n### Response:\n"
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     with torch.no_grad():
         outputs = model.generate(
@@ -40,10 +83,8 @@ def generate(instruction, temperature, max_tokens):
     code = result.split("### Response:")[-1].strip()
 
     # ── Always fix the topos line ─────────────────────────────
-    # Remove any existing broken topos line
     code = re.sub(r"\ntopos\s*=.*", "", code).strip()
 
-    # Find the class name
     match = re.search(r"class (\w+)\(Topo\)", code)
     if match:
         class_name = match.group(1)
@@ -55,13 +96,13 @@ examples = [
     ["Create a simple topology with 2 hosts and 1 switch.", 0.2, 512],
     ["Create a linear topology with 3 switches and 2 hosts per switch.", 0.2, 512],
     ["Build a tree topology with depth 2 and fanout 3.", 0.2, 512],
-    ["Create a topology with 3 routers using SoftRouter and TCLinks.", 0.2, 512],
+    ["Create a topology with 3 routers using TCLinks and bandwidth limits.", 0.2, 512],
 ]
 
 with gr.Blocks(title="Auto Mininet", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🌐 Auto Mininet
-    ### Convert plain English → Executable Mininet Python Script
+    ### Convert plain English → Executable Mininet Python Script  
     *Fine-tuned Qwen2.5-Coder-1.5B with LoRA*
     """)
 
@@ -87,13 +128,4 @@ with gr.Blocks(title="Auto Mininet", theme=gr.themes.Soft()) as demo:
     gr.Markdown("### 💡 How to run the generated script:")
     gr.Markdown("""
 ```bash
-    # Save the output to topo.py, then run:
-    sudo mn --custom topo.py --topo custom
-```
-    The topo name is always **custom** — same command every time.
-    """)
-
-    btn.click(fn=generate, inputs=[instruction, temperature, max_tokens], outputs=output)
-    gr.Examples(examples=examples, inputs=[instruction, temperature, max_tokens])
-
-demo.launch()
+sudo mn --custom topo.py --topo custom