---
library_name: transformers
pipeline_tag: text-generation
language:
  - en
datasets:
  - roneneldan/TinyStories
tags:
  - qraxai
  - gpt2-tokenizer
  - causal-lm
  - tiny-stories
  - from-scratch
license: mit
---

# QraXAi-Basic-45M

A small, decoder-only Transformer (~44.75M parameters) trained **from scratch** on
[TinyStories](https://huggingface.co/datasets/roneneldan/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

```python
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](https://arxiv.org/abs/2305.07759)).
- Tokenizer: GPT-2 BPE (`gpt2`).
- Built with PyTorch and Hugging Face Transformers.

## License

Released under the [MIT License](LICENSE). Note that the training dataset
([TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)) has its own terms;
check them if you plan to redistribute the data or use the model commercially.