Skip to content

Repository files navigation

Vashtra-0.6B

A 0.6B language model tuned for machine learning questions, small enough to run in a browser tab.

Demo · Model · Corpus · Project page

Weights are not published yet. This repository holds the full recipe. Weights and measured results follow when the training run finishes. There are no benchmark numbers here because none have been measured.

What it is

The base model is Qwen3-0.6B. It can discuss machine learning, but everything it knows about the field came from general web text. Vashtra continues training it on the ML literature itself, then fine-tunes it on machine learning questions that practitioners asked and answered.

Training

Stage Data What it teaches
A. Continued pretraining Full-text arXiv ML papers, plus 117k cs.LG abstracts Vocabulary, concepts, how a paper argues
B. Supervised fine-tuning ML, stats and CS Q&A, arXiv tasks, smoltalk2 mixture Answering, following instructions

Both stages are full-parameter. At 0.6B a full fine-tune beats LoRA for domain adaptation and still fits a free Colab T4, using about 10GB of 16GB.

Abstracts teach a model what papers claim; full text teaches it how the claim was argued. The corpus mixes both. About 22k general chat examples stay in the fine-tuning mixture, because tuning a 0.6B model purely on domain data reliably destroys its ability to hold a normal conversation.

Files

File Purpose
build_corpus.py Builds the corpus and pushes it to the Hub
train_vashtra.py Stage A, stage B, push the model
export_onnx.py ONNX export and quantisation for the browser demo
mask_check.py Verifies chat-template label masking
Vashtra_training.ipynb Colab notebook that runs the whole pipeline
space/ The static WebGPU demo hosted on Hugging Face Spaces
docs/ Project page published with GitHub Pages

Running it

Open Vashtra_training.ipynb in Colab, set the runtime to a T4 GPU, and run all. Everything pushes to your own Hugging Face repositories.

python build_corpus.py
python train_vashtra.py
python export_onnx.py

The corpus is around 90M tokens, more than one free Colab session will finish. Stage A takes a token budget and pushes checkpoints to the Hub, so --resume continues a run that was cut short instead of starting over.

python train_vashtra.py --cpt_tokens 30000000    # about two hours on a T4
python train_vashtra.py --resume                 # continue a dropped run

A detail worth copying

Qwen3's chat template inserts an empty <think></think> block into the final assistant turn. Tokenising message by message drops it, and you end up training on a format the model never sees at inference. mask_check.py verifies the encoder round-trips to apply_chat_template exactly. The failure is silent, so it is worth testing for.

Limits

It is a 0.6B model. It explains concepts reasonably and gets specifics wrong, particularly exact numbers, named papers and recent work. Do not cite it, and check anything that matters.

Licence

Apache-2.0, following the Qwen3-0.6B base. Corpus sources carry their own licences: full-text papers are CC-BY, CC-BY-SA or CC0, abstracts are AFL-3.0, StackExchange content is CC-BY-SA-4.0, and smoltalk2 is Apache-2.0. Anything derived from StackExchange stays CC-BY-SA-4.0.

About

A 0.6B language model tuned for machine learning questions, small enough to run in a browser tab

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages