Spaces:
Sleeping
Sleeping
| import torch | |
| import re | |
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from peft import PeftModel | |
| import ipaddress | |
| 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 | |
| base = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| torch_dtype=torch.float32, | |
| trust_remote_code=True | |
| ).to(device) | |
| model = PeftModel.from_pretrained(base, LORA_MODEL).to(device) | |
| model.eval() | |
| model.requires_grad_(False) | |
| print("✅ Model ready") | |
| SYSTEM_PROMPT = """ | |
| You are Auto Mininet. You ONLY output raw executable Mininet Python code. | |
| STRICT RULES: | |
| - Output ONLY raw Python code — no markdown, no backticks, no ```python, no ``` anywhere | |
| - NEVER use loops (no for, no while, no list comprehensions) | |
| - NEVER use f-strings or dynamic naming like f's{i}' or f'h{i}' | |
| - ALL switches must be named s1, s2, s3, s4... (canonical Mininet names) | |
| - ALL hosts must be named h1, h2, h3, h4... | |
| - Every switch and host must be declared on its own line explicitly | |
| - Every link must be on its own line explicitly | |
| - NEVER use root, child, node, sw or any non-standard switch names | |
| - NEVER answer questions unrelated to Mininet topology generation | |
| IP ADDRESS RULES: | |
| - If hosts connected to the same switch are assigned IPs from different subnets, prepend: | |
| # ERROR: Hosts on the same switch must be in the same subnet. Communication will fail. | |
| - When user specifies IP addresses, assign them using ip='x.x.x.x' inside addHost() | |
| - ALL hosts connected to the same switch MUST be in the same subnet | |
| - Example of CORRECT IPs on same switch: 10.0.0.1, 10.0.0.2, 10.0.0.3 (all /24) | |
| - Example of WRONG IPs on same switch: 10.0.0.1 and 192.168.1.1 (different subnets — will not communicate) | |
| - If user gives IPs from different subnets on the same switch, output this error at the top as a comment: | |
| # ERROR: Hosts on the same switch must be in the same subnet. Communication will fail. | |
| - If user gives IPs for only some hosts, assign default Mininet IPs to the rest (10.0.0.x) | |
| - Valid IP formats: 10.x.x.x, 192.168.x.x, 172.16.x.x | |
| - NEVER assign IPs outside of addHost() | |
| - NEVER hardcode subnet masks inside the script | |
| CORRECT EXAMPLE: | |
| from mininet.topo import Topo | |
| class MyTopo(Topo): | |
| def build(self): | |
| s1 = self.addSwitch('s1') | |
| s2 = self.addSwitch('s2') | |
| h1 = self.addHost('h1') | |
| h2 = self.addHost('h2') | |
| h3 = self.addHost('h3') | |
| self.addLink(s1, s2) | |
| self.addLink(s1, h1) | |
| self.addLink(s1, h2) | |
| self.addLink(s2, h3) | |
| topos = { 'custom': MyTopo } | |
| IP VALIDATION (STRICT — MUST FOLLOW): | |
| * A subnet is defined by the first 3 octets (assume /24) | |
| Example: | |
| 10.0.0.1 → subnet = 10.0.0 | |
| 10.0.0.2 → subnet = 10.0.0 ✅ SAME | |
| 192.168.1.1 → subnet = 192.168.1 ❌ DIFFERENT | |
| * ALL hosts connected to the SAME switch MUST have the SAME first 3 octets | |
| * BEFORE generating code: | |
| → Extract all IPs | |
| → Compare their first 3 octets | |
| * IF ANY mismatch exists: | |
| → You MUST add this EXACT line at the VERY TOP: | |
| # ERROR: Hosts on the same switch must be in the same subnet. Communication will fail. | |
| * DO NOT skip this check | |
| * DO NOT assume they can communicate | |
| --- | |
| ❌ WRONG CASE (MUST PRODUCE ERROR): | |
| User: | |
| "2 hosts 1 switch with IPs 10.0.0.1 and 192.168.1.1" | |
| Output MUST start with: | |
| # ERROR: Hosts on the same switch must be in the same subnet. Communication will fail. | |
| --- | |
| ✅ CORRECT CASE: | |
| IPs: | |
| 10.0.0.1 and 10.0.0.2 → OK (same subnet) | |
| --- | |
| ✅ EDGE CASE: | |
| If user gives: | |
| 10.0.1.1 and 10.0.2.1 | |
| → DIFFERENT subnet → ERROR | |
| --- | |
| * NEVER ignore subnet mismatch | |
| * NEVER silently generate code without error in this case | |
| CORRECT EXAMPLE WITH IP ADDRESSES: | |
| from mininet.topo import Topo | |
| class IPTopo(Topo): | |
| def build(self): | |
| s1 = self.addSwitch('s1') | |
| h1 = self.addHost('h1', ip='10.0.0.1') | |
| h2 = self.addHost('h2', ip='10.0.0.2') | |
| h3 = self.addHost('h3', ip='10.0.0.3') | |
| self.addLink(s1, h1) | |
| self.addLink(s1, h2) | |
| self.addLink(s1, h3) | |
| topos = { 'custom': IPTopo } | |
| CORRECT EXAMPLE WITH TWO SWITCHES AND DIFFERENT SUBNETS: | |
| from mininet.topo import Topo | |
| class TwoSubnet(Topo): | |
| def build(self): | |
| s1 = self.addSwitch('s1') | |
| s2 = self.addSwitch('s2') | |
| h1 = self.addHost('h1', ip='10.0.1.1') | |
| h2 = self.addHost('h2', ip='10.0.1.2') | |
| h3 = self.addHost('h3', ip='10.0.2.1') | |
| h4 = self.addHost('h4', ip='10.0.2.2') | |
| self.addLink(s1, s2) | |
| self.addLink(s1, h1) | |
| self.addLink(s1, h2) | |
| self.addLink(s2, h3) | |
| self.addLink(s2, h4) | |
| topos = { 'custom': TwoSubnet } | |
| - When user asks for IP addresses, use ip='x.x.x.x' inside addHost() | |
| - When user asks for IPs from different subnets on the same switch, add error comment at top | |
| - NEVER assign IPs outside of addHost() | |
| WRONG — NEVER do this: | |
| ````python <- NEVER add this | |
| for i in range(3): <- NEVER use loops | |
| self.addSwitch(f's{i}') <- NEVER use f-strings | |
| root = self.addSwitch('root') <- NEVER use non-standard names | |
| ``` <- NEVER add this | |
| """ | |
| NETWORK_KEYWORDS = [ | |
| "topology", "mininet", "switch", "host", "router", | |
| "link", "tree", "linear", "fat-tree", "tc", "bandwidth", "delay", | |
| "network", "switches", "hosts", "routers", "create", "build", | |
| "generate", "design", "connect", "nodes", "topo", "hub", | |
| "ring", "mesh", "star", "hierarchical", "core", "edge", | |
| "aggregation", "packet", "port", "interface", "vlan" | |
| ] | |
| # ── Subnet conflict detection ────────────────────────── | |
| def check_subnet_conflict(code): | |
| host_ips = {} | |
| for match in re.finditer(r"addHost\(\s*['\"]([^'\"]+)['\"]\s*,\s*ip\s*=\s*['\"]([^'\"]+)['\"]", code): | |
| host_ips[match.group(1)] = match.group(2) | |
| if len(host_ips) < 2: | |
| return None | |
| switch_hosts = {} | |
| links = re.findall(r"self\.addLink\(\s*([^,]+)\s*,\s*([^,)]+)", code) | |
| for n1_raw, n2_raw in links: | |
| n1 = n1_raw.replace("self.", "").strip(" '\"") | |
| n2 = n2_raw.replace("self.", "").strip(" '\"") | |
| if n1.startswith('s') and n2 in host_ips: | |
| switch_hosts.setdefault(n1, []).append(n2) | |
| elif n2.startswith('s') and n1 in host_ips: | |
| switch_hosts.setdefault(n2, []).append(n1) | |
| for switch, connected_hosts in switch_hosts.items(): | |
| if len(connected_hosts) > 1: | |
| networks = set() | |
| for h in connected_hosts: | |
| ip_str = host_ips[h] | |
| if '/' not in ip_str: | |
| ip_str += '/24' | |
| try: | |
| networks.add(ipaddress.ip_network(ip_str, strict=False)) | |
| except ValueError: | |
| pass | |
| if len(networks) > 1: | |
| return f"# ⚠️ SUBNET CONFLICT ERROR: Hosts on switch '{switch}' are in different subnets: {[str(n) for n in networks]}\n# Hosts in different subnets on the same switch cannot communicate without a router.\n\n" | |
| return None | |
| def is_network_query(text: str) -> bool: | |
| text = text.lower() | |
| reject_keywords = [ | |
| "food", "recipe", "weather", "movie", "music", "sport", | |
| "hello", "hi", "how are you", "joke", "story", "poem", | |
| "capital", "president", "history", "math", "calculate" | |
| ] | |
| if any(k in text for k in reject_keywords): | |
| return False | |
| if any(k in text for k in NETWORK_KEYWORDS): | |
| return True | |
| topology_phrases = ["create a", "build a", "make a", "generate a", "design a"] | |
| if any(p in text for p in topology_phrases): | |
| return True | |
| return False | |
| def generate(instruction, temperature, max_tokens): | |
| if not instruction.strip(): | |
| return "⚠️ Please enter a topology description." | |
| 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( | |
| **inputs, | |
| max_new_tokens=int(max_tokens), | |
| temperature=float(temperature), | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| result = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| code = result.split("### Response:")[-1].strip() | |
| # Strip markdown backticks if model adds them | |
| code = re.sub(r"^```python\s*", "", code) | |
| code = re.sub(r"^```\s*", "", code) | |
| code = re.sub(r"\s*```$", "", code) | |
| code = code.strip() | |
| # Remove special model tokens | |
| code = re.sub(r"<\|.*?\|>", "", code).strip() | |
| # Remove any existing broken topos line and re-add clean one | |
| code = re.sub(r"\ntopos\s*=.*", "", code).strip() | |
| match = re.search(r"class (\w+)\(Topo\)", code) | |
| if match: | |
| class_name = match.group(1) | |
| code += f"\n\ntopos = {{ 'custom': {class_name} }}" | |
| error = check_subnet_conflict(code) | |
| if error: | |
| code = error | |
| return code | |
| 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 TCLinks and bandwidth limits.", 0.2, 512], | |
| ["Create a topology with 2 hosts and 1 switch, assign IP 10.0.0.1 to h1 and 10.0.0.2 to h2.", 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 | |
| *Fine-tuned Qwen2.5-Coder-1.5B with LoRA* | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| instruction = gr.Textbox( | |
| label="📝 Describe your network topology", | |
| placeholder="e.g. Create a topology with 4 hosts and 2 switches...", | |
| lines=5 | |
| ) | |
| with gr.Row(): | |
| temperature = gr.Slider(0.1, 1.0, value=0.2, step=0.1, label="🎲 Creativity") | |
| max_tokens = gr.Slider(128, 1024, value=512, step=64, label="📏 Max length") | |
| btn = gr.Button("⚡ Generate Mininet Script", variant="primary", size="lg") | |
| with gr.Column(): | |
| output = gr.Code( | |
| label="🐍 Generated Python Script", | |
| language="python", | |
| lines=20 | |
| ) | |
| gr.Markdown("### 💡 How to run the generated script:") | |
| gr.Markdown(""" | |
| ```bash | |
| 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() |