feat: Add universal PyTorch/Transformers backend for Linux and Google Colab

ControlAI_Colab_Demo.ipynb

@@ -7,13 +7,10 @@
     "# ControlAI: Open-Source Safety-Critical AI Agent for Control Systems\n",
     "\n",
     "[![GitHub Repo](https://img.shields.io/badge/GitHub-atakankahya%2Fcontrolai--agent-blue?logo=github)](https://github.com/atakankahya/controlai-agent)\n",
-    "[![Hugging Face Model](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-ControlAI--Agent-yellow)](https://huggingface.co/atakankahya/ControlAI-Agent)\n",
+    "[![Hugging Face Model](https://img.shields.io/badge/Hugging%20Face-ControlAI--Agent-blue)](https://huggingface.co/atakankahya/ControlAI-Agent)\n",
     "[![License](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)\n",
     "\n",
-    "**ControlAI** is an open-source, safety-oriented AI agent specialized in Control Systems Engineering, Applied Mathematics, and Dynamical Systems Simulation. Unlike black-box general LLMs that hallucinate stability margins and invert singular matrices, ControlAI enforces:\n",
-    "1. **Deterministic Python Sandbox:** Exact SciPy/Control/CVXPY ODE simulations and plot generation.\n",
-    "2. **4-Stage Mathematical Proof Standard:** Formal Lyapunov, PBH rank, and Bode sensitivity trade-offs.\n",
-    "3. **Zero-Hallucination Guarantees:** Rigorous numerical verification with real engineering bounds.\n",
+    "**ControlAI** is an open-source AI agent specialized in Control Systems Engineering, Applied Mathematics, and Dynamical Systems Simulation.\n",
     "\n",
     "---"
    ]
@@ -22,7 +19,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Step 1: Clone Repository & Install Dependencies"
+    "### Step 1: Clone Repository & Install PyTorch/Transformers Dependencies"
    ]
   },
   {
@@ -35,15 +32,15 @@
     "!git clone https://github.com/atakankahya/controlai-agent.git\n",
     "%cd controlai-agent\n",
     "\n",
-    "# Install core scientific control engineering and web dependencies\n",
-    "!pip install -q scipy numpy matplotlib control cvxpy rich pyngrok uvicorn fastapi"
+    "# Install PyTorch, Transformers, and Scientific Control Libraries\n",
+    "!pip install -q torch transformers accelerate scipy control cvxpy rich pyngrok uvicorn fastapi sentencepiece protobuf python-multipart"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Step 2: Quick CLI Interactive Test (Terminal Mode in Colab)"
+    "### Step 2: Quick CLI Test in Terminal Mode"
    ]
   },
   {
@@ -52,16 +49,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Test single control engineering query\n",
-    "!python cli.py \"Design an LQR controller for A=[[0, 1], [-2, -3]], B=[[0], [1]], Q=diag([10, 1]), R=1 and simulate step response.\""
+    "# Run single control engineering query using PyTorch/Transformers\n",
+    "!python cli.py --model \"Qwen/Qwen2.5-3B-Instruct\" \"Design an LQR controller for A=[[0, 1], [-2, -3]], B=[[0], [1]], Q=diag([10, 1]), R=1 and simulate step response.\""
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Step 3: Launch Live Public Web Interface (Cloud Demo Link)\n",
-    "Run the cell below to generate a public URL and open the full ControlAI Web Console in your browser!"
+    "### Step 3: Launch Live Public Web Interface (Interactive Playground Link)\n",
+    "Run the cell below to launch the web console and get your public web URL!"
    ]
   },
   {
@@ -70,10 +67,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "import subprocess\n",
     "import time\n",
     "\n",
-    "# Start the FastAPI server in background\n",
+    "# Launch FastAPI server\n",
     "server_process = subprocess.Popen([\"python\", \"-m\", \"uvicorn\", \"app:app\", \"--host\", \"0.0.0.0\", \"--port\", \"8000\"])\n",
     "time.sleep(3)\n",
     "\n",

controlai_agent/orchestrator.py

@@ -8,8 +8,16 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Generator
 
-from mlx_lm import generate, load, stream_generate
-from transformers import AutoTokenizer
+import os
+import sys
+
+try:
+    from mlx_lm import generate as mlx_generate, load as mlx_load
+    HAS_MLX = True
+except ImportError:
+    HAS_MLX = False
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from controlai_agent.prompts import CONTROLAI_SYSTEM_PROMPT
 from controlai_agent.registry import ToolRegistry, registry
@@ -78,8 +86,7 @@ def _extract_tool_calls(text: str) -> tuple[list[dict[str, Any]], str]:
                 calls.append(obj)
 
     # 3. Raw JSON object containing "name" and "arguments" / "parameters"
-    if not calls and '"name"' in text or "'name'" in text:
-        # Match starting from first '{' to last matching '}'
+    if not calls and ('"name"' in text or "'name'" in text):
         match = re.search(r"(\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*\})", text)
         if match:
             obj = parse_flexible_json(match.group(1))
@@ -95,7 +102,7 @@ def _extract_tool_calls(text: str) -> tuple[list[dict[str, Any]], str]:
 
 
 class ControlAIAgent:
-    """Offline Control Engineering Agent using Qwen3-4B on MLX with deterministic tool calling and streaming."""
+    """Universal Control Engineering Agent supporting MLX (Apple Silicon) and PyTorch/Transformers (Linux/CUDA)."""
 
     def __init__(
         self,
@@ -109,12 +116,31 @@ class ControlAIAgent:
         self.registry = tool_registry
         self.max_tool_steps = max_tool_steps
 
-        # Load MLX model and HF tokenizer
-        if adapter_path:
-            self.model, self.mlx_tokenizer = load(model_path, adapter_path=adapter_path)
+        # Detect platform & backend
+        self.is_mlx = HAS_MLX and not model_path.startswith("Qwen/") and not os.environ.get("FORCE_TRANSFORMERS")
+
+        if self.is_mlx:
+            if adapter_path:
+                self.model, self.mlx_tokenizer = mlx_load(model_path, adapter_path=adapter_path)
+            else:
+                self.model, self.mlx_tokenizer = mlx_load(model_path)
+            self.hf_tokenizer = AutoTokenizer.from_pretrained(model_path)
         else:
-            self.model, self.mlx_tokenizer = load(model_path)
-        self.hf_tokenizer = AutoTokenizer.from_pretrained(model_path)
+            # Universal PyTorch / Transformers fallback on Linux, Colab, HuggingFace, CUDA
+            hf_id = "Qwen/Qwen2.5-3B-Instruct" if "mlx" in model_path else model_path
+            self.hf_tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                hf_id,
+                torch_dtype="auto",
+                device_map="auto",
+                trust_remote_code=True,
+            )
+            if adapter_path and Path(adapter_path).exists():
+                try:
+                    from peft import PeftModel
+                    self.model = PeftModel.from_pretrained(self.model, adapter_path)
+                except Exception:
+                    pass
 
         # Initialize local offline RAG index
         try:
@@ -122,6 +148,29 @@ class ControlAIAgent:
         except Exception:
             self.rag_index = None
 
+    def _generate(self, prompt: str, max_tokens: int = 2000) -> str:
+        """Universal text generation handling both MLX and PyTorch backends."""
+        if self.is_mlx:
+            return mlx_generate(
+                self.model,
+                self.mlx_tokenizer,
+                prompt=prompt,
+                max_tokens=max_tokens,
+                verbose=False,
+            ).strip()
+        else:
+            import torch
+            inputs = self.hf_tokenizer(prompt, return_tensors="pt").to(self.model.device)
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=max_tokens,
+                    do_sample=False,
+                    pad_token_id=self.hf_tokenizer.eos_token_id,
+                )
+            new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+            return self.hf_tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+
     def _get_grounded_instruction(self, user_prompt: str, base_instruction: str) -> str:
         """Retrieve relevant textbook theorems and inject grounding into system instructions."""
         if not self.rag_index or not self.rag_index.chunks:
@@ -181,13 +230,7 @@ class ControlAIAgent:
                 add_generation_prompt=True,
             )
 
-            model_output = generate(
-                self.model,
-                self.mlx_tokenizer,
-                prompt=rendered_prompt,
-                max_tokens=max_tokens_per_step,
-                verbose=False,
-            ).strip()
+            model_output = self._generate(rendered_prompt, max_tokens=max_tokens_per_step)
 
             tool_calls, pre_text = _extract_tool_calls(model_output)
 
@@ -240,13 +283,7 @@ class ControlAIAgent:
             tokenize=False,
             add_generation_prompt=True,
         )
-        final_output = generate(
-            self.model,
-            self.mlx_tokenizer,
-            prompt=forced_prompt,
-            max_tokens=max_tokens_per_step,
-            verbose=False,
-        ).strip()
+        final_output = self._generate(forced_prompt, max_tokens=max_tokens_per_step)
 
         _, clean_final = _extract_tool_calls(final_output)
         return AgentResult(
@@ -299,13 +336,7 @@ class ControlAIAgent:
                 add_generation_prompt=True,
             )
 
-            model_output = generate(
-                self.model,
-                self.mlx_tokenizer,
-                prompt=rendered_prompt,
-                max_tokens=max_tokens_per_step,
-                verbose=False,
-            ).strip()
+            model_output = self._generate(rendered_prompt, max_tokens=max_tokens_per_step)
 
             tool_calls, pre_text = _extract_tool_calls(model_output)
 
@@ -396,13 +427,7 @@ class ControlAIAgent:
             add_generation_prompt=True,
         )
 
-        final_output = generate(
-            self.model,
-            self.mlx_tokenizer,
-            prompt=forced_prompt,
-            max_tokens=max_tokens_per_step,
-            verbose=False,
-        ).strip()
+        final_output = self._generate(forced_prompt, max_tokens=max_tokens_per_step)
 
         # If model generated another tool call during final synthesis, execute it!
         synth_tool_calls, clean_synth = _extract_tool_calls(final_output)
@@ -421,7 +446,7 @@ class ControlAIAgent:
 
             # Re-generate synthesis after executing the tool
             re_prompt = self.hf_tokenizer.apply_chat_template(messages, tools=None, tokenize=False, add_generation_prompt=True)
-            final_output = generate(self.model, self.mlx_tokenizer, prompt=re_prompt, max_tokens=max_tokens_per_step, verbose=False).strip()
+            final_output = self._generate(re_prompt, max_tokens=max_tokens_per_step)
             _, clean_synth = _extract_tool_calls(final_output)
 
         final_text = clean_synth or final_output