balalida/gpt-22M-owt

CS336 openwebtext causal language model

Four-layer decoder-only Transformer trained for Stanford CS336 Assignment 1. It uses a byte-level BPE tokenizer, context length 256, width 512, 16 attention heads, SwiGLU width 1344, RMSNorm, and RoPE.

• —Validation loss: 4.140625 per token
• —Training tokens represented by this checkpoint lineage: 647,168,000
• —Parameters are tied between the token embedding and LM head.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("PATH_OR_REPO", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("PATH_OR_REPO", trust_remote_code=True)
inputs = tokenizer("The United States", return_tensors="pt")
output = model.generate(**inputs, max_new_tokens=128, temperature=0.8, top_p=0.9, do_sample=True)
print(tokenizer.decode(output[0], skip_special_tokens=False))

This is a small educational model and may hallucinate, repeat, or produce incorrect information.