QraXAi-Basic-45M

A small, decoder-only Transformer (~44.75M parameters) trained from scratch on TinyStories with the GPT-2 BPE tokenizer (50,257 tokens). QraXAi is a hand-written PyTorch GPT implementation (model.py / configuration_qraxai.py, shipped with this repo), not a fine-tune of GPT-2 — only the tokenizer is shared with GPT-2.

This is an experimental research model, intended for learning/demo purposes.

Model details

Architecture	Decoder-only Transformer (GPT-style), pre-norm
Parameters	44,751,872 (~44.75M), all trainable
Layers	24
Hidden size	256
Attention heads	8 (head dim 32)
Feed-forward	4× hidden, GELU
Context length	256 tokens (hard limit)
Vocabulary	50,257 (GPT-2 BPE)
Position encoding	Learned absolute embeddings
Normalization	LayerNorm
Weight tying	No (lm_head is separate)
KV cache	No — generation recomputes the full context at every step
Weights	fp32, 179 MB (model.safetensors)
Special tokens	bos = eos = <|endoftext|> (id 50256)
Custom code	Yes — requires trust_remote_code=True

Parameter breakdown: token embeddings 12.87M + position embeddings 0.07M + 24 × 0.79M transformer blocks (18.95M) + final norm + untied lm_head 12.87M.

Quick start

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "coderian/QraXAi-Basic-45M"

tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
    dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
).to("cuda" if torch.cuda.is_available() else "cpu").eval()

prompt = "Once upon a time, there was a little girl named Lily"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=244,          # prompt + new tokens must stay <= 256
        do_sample=True,
        temperature=0.8,
        top_k=50,
        top_p=0.95,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

print(tokenizer.decode(output[0], skip_special_tokens=True))

Notes for generation:

• Context is a hard limit of 256 tokens. The custom forward raises ValueError if the sequence gets longer, so keep len(prompt) + max_new_tokens <= 256.
• The model has no KV cache: each new token re-runs the full context, so generation cost grows quickly with sequence length.
• The custom forward does not use an attention mask for padding. Generate one prompt at a time instead of batching.
• Generation usually stops at <|endoftext|>, but because the training data is a continuous stream of stories, the model sometimes starts a new story instead.

Example output

Prompt: Once upon a time, there was a little girl named Lily (temperature=0.8, top_k=50, top_p=0.95, repetition_penalty=1.1, max_new_tokens=244):

Once upon a time, there was a little girl named Lily who loved to play in the big, green field. One day, she found a shiny stone on top of her backyard. She picked it up and showed it to her mom.

"Look mommy, I found a pretty mineral!" said Lily excitedly. "It's very pretty!"

Her mom smiled and said, "That's right, sweetie. It'll make sure you touch it. But remember, be careful with it because you might find something else inside."

Lily nodded her head and kept playing with the jewel until she noticed that the box had fallen into a hole. She felt sad for her mom, but then remembered what her mom said about when something is hurt.

The next day, Lily went back to the park and saw that the unknown stone was broken. She asked her mom if they could try and fix it. Her mom told her that it's okay to ask for help and that sometimes you can't use it without asking permission. So, Lily listened to her mom and never touched the stone again. Once upon a time, there was a boy named Timmy. He loved to play with his toy car. One day, he went to

(The last line is a new story the model began, cut off by the token limit.)

Training

Dataset	roneneldan/TinyStories (train split, streaming), first 130,000 stories
Data size	115.7M characters → 28.76M tokens → ~112,350 training blocks
Objective	Next-token prediction (causal LM), cross-entropy
Epochs	1 (~7,000 optimizer steps)
Batch size	16
Block size	256
Optimizer	AdamW, lr 3e-4
Gradient clipping	1.0
Mixed precision	bf16 (on CUDA)
Seed	42

Files in this repository

File	Description
config.json	Model config (GPTConfig + auto_map for remote code)
model.safetensors	fp32 weights (179 MB)
model.py	QraXAiForCausalLM — custom modeling code
configuration_qraxai.py	GPTConfig — custom configuration
tokenizer.json, tokenizer_config.json	GPT-2 BPE tokenizer
generation_config.json	Default generation settings
LICENSE	MIT License

Limitations

• English only; trained exclusively on synthetic children's stories (TinyStories), so it knows little about the real world.
• Trained for a single epoch: grammar is mostly coherent, but content can be repetitive, inconsistent or nonsensical.
• Hard 256-token context; no sliding window, so long inputs must be truncated.
• Not instruction-tuned — it does not follow instructions and is not a chat model.
• No safety filtering or alignment of any kind. Do not use in production or for user-facing applications.
• TinyStories is synthetic data; nothing prevents the model from producing odd or inappropriate continuations.

Acknowledgements

• Dataset and idea: TinyStories: How Small Can Language Models Be and Still Speak Coherent English? — Ronen Eldan and Yuanzhi Li, 2023 (arXiv:2305.07759).
• Tokenizer: GPT-2 BPE (gpt2).
• Built with PyTorch and Hugging Face Transformers.

License

Released under the MIT License. Note that the training dataset (TinyStories) has its own terms; check them if you plan to redistribute the data or use the model commercially.