{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ⚛️ Q-TensorFormer: 1-Click Interactive Google Colab Quickstart\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Premchan369/Q-TensorFormer/blob/main/notebooks/quickstart.ipynb)\n",
    "[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Premchan369%2FQ--TensorFormer-yellow)](https://huggingface.co/Premchan369/Q-TensorFormer)\n",
    "\n",
    "Welcome to **Q-TensorFormer**, an arXiv-grade, hardware-aware Transformer architecture that dynamically allocates model capacity, tensor rank ($r \\in [4, 16]$), and KV cache precision on a per-token basis using a closed-loop PID controller.\n",
    "\n",
    "In this notebook, you will:\n",
    "1. Load Q-TensorFormer directly with `trust_remote_code=True`.\n",
    "2. Observe token-by-token dynamic rank shifting and live energy telemetry in action.\n",
    "3. Learn how to compress any dense Hugging Face model (LLaMA, Qwen) into a low-rank Q-TensorFormer model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q \"transformers>=4.38.0\" \"torch>=2.0.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Load Q-TensorFormer Directly via Hugging Face AutoModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Load turnkey model from Hugging Face Hub\n",
    "model_id = \"Premchan369/Q-TensorFormer\"\n",
    "config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to(device)\n",
    "model.eval()\n",
    "\n",
    "print(f\"✓ Successfully loaded Q-TensorFormer ({sum(p.numel() for p in model.parameters()):,} parameters)!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Autoregressive Text Generation with Live Dynamic Telemetry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example prompt\n",
    "prompt = \"Quantum tensor network architectures optimize energy by\"\n",
    "print(f\"Input Prompt: {prompt}\\n\")\n",
    "\n",
    "# Tokenize prompt\n",
    "tokens = [ord(c) % config.vocab_size for c in prompt]\n",
    "input_ids = torch.tensor([tokens], dtype=torch.long, device=device)\n",
    "\n",
    "# Generate with live closed-loop PID rank tracking\n",
    "curr_ids = input_ids\n",
    "past_key_values = None\n",
    "\n",
    "print(\"--- Streaming Generated Output with Live Tensor-Train Rank ---\")\n",
    "for step in range(32):\n",
    "    with torch.no_grad():\n",
    "        outputs = model(curr_ids[:, -1:], past_key_values=past_key_values, use_cache=True)\n",
    "        logits = outputs.logits[:, -1, :]\n",
    "        past_key_values = outputs.past_key_values\n",
    "        next_tok = logits.argmax(dim=-1, keepdim=True)\n",
    "    \n",
    "    curr_ids = torch.cat([curr_ids, next_tok], dim=1)\n",
    "    tok_id = int(next_tok[0, 0].item())\n",
    "    char = chr(tok_id) if 32 <= tok_id < 127 else \" \"\n",
    "    \n",
    "    # Query live PID controller rank\n",
    "    pid_ctrl = model.model.pid_controller\n",
    "    active_rank = int(pid_ctrl.rank_levels[int(pid_ctrl.active_rank_idx.item())])\n",
    "    \n",
    "    print(f\"{char}\", end=\"\", flush=True)\n",
    "\n",
    "print(\"\\n\\n✓ Generation Complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Drop-In LLM Compressor\n",
    "\n",
    "You can also use Q-TensorFormer to compress existing popular dense models (like LLaMA-3.2-1B, Qwen-2.5, or Mistral) down by 50%:\n",
    "\n",
    "```bash\n",
    "!python qtensorformer_compress.py \\\n",
    "    --model-name-or-path meta-llama/Llama-3.2-1B \\\n",
    "    --output-dir ./Llama-3.2-1B-QTensor \\\n",
    "    --tt-rank 16\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}