File size: 3,857 Bytes
d073da3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | ---
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.
|