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

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

Files in this repository

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.