trainai-essays

Exported language model

A small base language model trained from scratch with TrainAI. It completes text; it is not an instruction-following assistant. Ask it a question and it will continue the question rather than answer it.

What it is

• Parameters: 5.3M
• Layers: 4, heads: 4 (key/value heads: 4)
• Model dimension: 256
• Context length: 256 tokens
• Vocabulary: 8192 tokens, byte-level BPE trained on the same corpus
• Weights: float32
• Trained for 1862 steps
• Best validation loss: 2.8390 (at step 1862)

This is a LlamaForCausalLM -- the architecture TrainAI trains is structurally a Llama (RMSNorm, rotary embeddings, SwiGLU, grouped-query attention, no biases), so the weights are the same tensors under HuggingFace's names.

Using it

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("./")
model = AutoModelForCausalLM.from_pretrained("./")

ids = tokenizer("Once upon a time", return_tensors="pt")
print(tokenizer.decode(model.generate(**ids, max_new_tokens=50)[0]))

Provenance

Exported from step-0001862.pt (best) of the run at runs/essays-run. Tokenizer fingerprint a31e62e0d4fb4e892c9c9c80f4f2f940a75a818cba9f7c929de0bae1cf6a56f5 -- the tokenizer in this directory is the one the model was trained with, and no other tokenizer will produce sensible output from these weights.