OpenLanguage Model.

import sys, os, torch, urllib.request
from torch.utils.data import DataLoader
from tempfile import TemporaryDirectory
from olm import Dataset, HFTokenizer, Trainer, LM

with TemporaryDirectory() as tmp:
    urllib.request.urlretrieve("https://github.com/.../input.txt",
                              os.path.join(tmp, "i.txt"))
    tokenizer, device = HFTokenizer("gpt2"), "cuda" if torch.cuda.is_available() else "cpu"

    # Define Model
    model = LM(tokenizer.vocab_size, 64, 4, 2, 33)
    optimizer = torch.optim.AdamW(model.parameters(), 3e-4)

    # Data & Training
    dataset = Dataset(tmp, tokenizer, 32)
    dataloader = DataLoader(dataset, 4)
    trainer = Trainer(model, optimizer, dataloader, device, 32, use_amp=False)

    losses = trainer.train(1, 10, 100)
    print(f"S:{losses[0]:.4f} E:{losses[-1]:.4f} OK:{losses[-1]<losses[0]}")

import torch
import torch.nn.functional as F
from olm.nn.activations.base import ActivationBase

class SwiGLU(ActivationBase):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        value, gate = x.chunk(2, dim=-1)
        return value * F.silu(gate)

Llama3Block = Block([
    Residual(Block([
        RMSNorm(embed_dim, eps=1e-5),
        GroupedQueryAttention(
            embed_dim, num_heads, num_kv_heads, max_seq_len,
            dropout=dropout, rope_theta=rope_theta, use_bias=False
        )
    ])),
    Residual(Block([
        RMSNorm(embed_dim, eps=1e-5),
        SwiGLUFFN(embed_dim, hidden_dim=intermediate_size,
                  dropout=dropout, bias=False)
    ]))
])

Llama3Model = Block([
    Embedding(vocab_size, embed_dim),
    Repeat(lambda: Llama3Block(
        embed_dim, intermediate_size, num_heads, num_kv_heads,
        max_seq_len, dropout, rope_theta
    ), num_layers),
    RMSNorm(embed_dim, eps=1e-5),
    Linear(embed_dim, vocab_size, bias=False)
])