File size: 4,737 Bytes
3d4f2e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | {
"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
}
|