{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vortex Alpha\n",
    "\n",
    "A compact, experimental 174.9M-parameter language model. The final public name is still undecided.\n",
    "\n",
    "This notebook loads the public Hugging Face checkpoint with the standard Transformers API. Vortex is a research preview: expect factual, arithmetic, repetition, and long-context errors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the small runtime\n",
    "\n",
    "On Colab, select a GPU runtime when available. The model is small enough to fit comfortably in a typical Colab GPU, although this reference implementation does not use a KV cache."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip -q install -U \"transformers>=4.45\" sentencepiece safetensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Load Vortex from the Hub\n",
    "\n",
    "`trust_remote_code=True` is required because Vortex has a custom GQA + QK-Norm implementation. The repository contains the configuration, model, tokenizer, and generation code used here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "REPO = \"North-ML1/vortex-alpha\"\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    dtype = torch.float32\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    REPO, trust_remote_code=True, torch_dtype=dtype\n",
    ").to(device).eval()\n",
    "\n",
    "print(\"device:\", device)\n",
    "print(\"dtype:\", next(model.parameters()).dtype)\n",
    "print(\"parameters:\", model.num_parameters())\n",
    "print(\"tokenizer vocabulary:\", tokenizer.vocab_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Ask a question with the built-in chat template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask(question: str, max_new_tokens: int = 96) -> str:\n",
    "    messages = [{\"role\": \"user\", \"content\": question}]\n",
    "    prompt = tokenizer.apply_chat_template(\n",
    "        messages, tokenize=False, add_generation_prompt=True\n",
    "    )\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(device)\n",
    "    with torch.inference_mode():\n",
    "        generated = model.generate(\n",
    "            **inputs,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "            do_sample=False,\n",
    "            pad_token_id=tokenizer.pad_token_id,\n",
    "            eos_token_id=tokenizer.eos_token_id,\n",
    "        )\n",
    "    new_tokens = generated[0, inputs[\"input_ids\"].shape[1]:]\n",
    "    return tokenizer.decode(new_tokens, skip_special_tokens=True)\n",
    "\n",
    "print(ask(\"Explain why the sky appears blue in two short paragraphs.\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "questions = [\n",
    "    \"Solve 3x + 5 = 20 and show the steps.\",\n",
    "    \"Write a short Python function that returns the largest number in a list.\",\n",
    "    \"Summarize: The museum opens at 9, closes at 5, and admission is free on Sunday.\",\n",
    "]\n",
    "for question in questions:\n",
    "    print(f\"\\nUSER: {question}\\nVORTEX: {ask(question)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notes\n",
    "\n",
    "- `model.safetensors` is the experimental instruction/tool-format preview.\n",
    "- `base_model.safetensors` is the corresponding pretrained base; the notebook loads the instruction preview by default.\n",
    "- The model may emit a `CALL {json}` tool request, but this notebook does not provide external tools.\n",
    "- Do not rely on Vortex for medical, legal, financial, or other high-stakes decisions."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}