t-nano-16.6.26

t-1: 100M parameters, 24M active, 1B training tokens

This is a small PyTorch LM that was trained on 500M tokens from FineWeb EDU, and 500M tokens from Wikipedia Monthly. It was only trained on English text data.

Inference

Run the model with this simple script:

import torch
from huggingface_hub import hf_hub_download

from model import MoELMConfig, MoELanguageModel
from tokenizer_utils import EOS_TOKEN, decode, encode, load_tokenizer, token_id

repo = "owenqwenllmwine/t-nano"
device = "cuda" if torch.cuda.is_available() else "cpu"

ckpt_path = hf_hub_download(repo, "final.pt")
tok_path = hf_hub_download(repo, "tokenizer.json")

tokenizer = load_tokenizer(tok_path)

ckpt = torch.load(ckpt_path, map_location=device)
model = MoELanguageModel(MoELMConfig.from_dict(ckpt["config"])).to(device)
model.load_state_dict(ckpt["model"])
model.eval()

prompt = "The history of artificial intelligence begins"
x = torch.tensor([encode(tokenizer, prompt, add_bos=True)], device=device)

with torch.no_grad():
    y = model.generate(
        x,
        max_new_tokens=200,
        eos_token_id=token_id(tokenizer, EOS_TOKEN),
        temperature=0.8,
        top_k=50,
        top_p=0.95,
    )

print(decode(tokenizer, y[0].tolist(), skip_special_tokens=True))