README.md

---
license: apache-2.0
language:
  - en
library_name: transformers
pipeline_tag: text-generation
datasets:
  - exnivo/tinybrain-pretrain-corpus-2b
tags:
  - causal-lm
  - gpt2
  - small-language-model
  - tinybrain
---

# 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

```python
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`](https://huggingface.co/datasets/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.