QraXAi
QraXAi is a small decoder-only transformer language model (~32M parameters) trained
from scratch on English text. It was built as a minimal, readable GPT-style
implementation with a custom architecture (see model.py).
This is a base model: it continues text, it does not follow instructions.
Model details
| Architecture | GPT-style causal decoder-only transformer |
| Parameters | 32.1M |
| Layers | 8 |
| Hidden size | 256 |
| Attention heads | 8 (head dim 32) |
| Context length | 256 tokens |
| Vocabulary | 50,257 |
| Tokenizer | GPT-2 BPE (GPT2TokenizerFast) |
| Positional encoding | Learned absolute embeddings |
| Normalization | Pre-LayerNorm |
| MLP | GELU, 4x expansion |
| Weight dtype | float32 |
| Auto class | AutoModelForCausalLM (custom code, trust_remote_code=True) |
Uses
Direct use: text completion — give an English prompt, get a continuation.
Out of scope: chat / instruction following (not fine-tuned), factual question answering, non-English text, prompts longer than 256 tokens (prompt + generated tokens combined).
Limitations
- Undertrained: only 1 epoch (~61M tokens) was trained. Training loss went from 10.85 (random) to 6.02 (perplexity ~410), so the output is locally plausible but often incoherent.
- Short context: the model was trained with 256-token blocks and cannot attend beyond that.
- No KV cache: generation recomputes the full context at every step, so sampling is slow.
- No attention mask / padding support: generate a single sequence at a time; padded batches would produce wrong results.
- English only; may reproduce biases present in the training corpus.
How to use
import torch
from transformers import AutoModelForCausalLM, GPT2TokenizerFast
model_id = "coderian/qraxai"
tokenizer = GPT2TokenizerFast.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
trust_remote_code=True, # custom architecture, code ships with the model
).eval()
prompt = "The purpose of this experiment is"
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=64,
do_sample=True,
temperature=0.8,
top_k=50,
pad_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(outputs[0]))
Training
Data: exnivo/tinybrain-pretrain-corpus-2b
(first 75,000 rows) — English educational / web / code text, ~284 MB, ~61M GPT-2 tokens.
Procedure
| Setting | Value |
|---|---|
| Epochs | 1 |
| Blocks / batch | 256 tokens / 16 |
| Tokens per step | 4,096 |
| Optimizer steps | 15,584 |
| Optimizer | AdamW |
| Learning rate | 3e-4 (constant, no warmup) |
| Gradient clipping | 1.0 |
| Precision | bf16 autocast, fp32 master weights |
| Hardware | NVIDIA RTX 4060 Laptop (8 GB) |
| Time | ~21 min / epoch |
| Final loss | 6.02 (from 10.85 at initialization) |
No validation split or benchmark evaluation was run — use the loss numbers as a rough signal of quality only.
Files
config.json # GPTConfig + auto_map
model.safetensors # 32.1M parameters, fp32
model.py # QraXAiForCausalLM
configuration_qraxai.py # GPTConfig
tokenizer.json # GPT-2 BPE
tokenizer_config.json
generation_config.json
License
Apache-2.0 (see LICENSE). The tokenizer files come from OpenAI's GPT-2 release
(MIT license). The training corpus is credited above; check its terms before
redistributing derivatives.