OpenLanguage Model.

import torch

from olm.nn.blocks import LM
from olm.train import Trainer
from olm.data.tokenization import HFTokenizer
from olm.data.datasets import FineWebEduDataset, DataLoader

tok = HFTokenizer("gpt2")
model = LM(
    tok.vocab_size,
    embed_dim=640,
    num_heads=10,
    num_layers=12,
    max_seq_len=1024,
    ff_multiplier=2.75,
)

dataset = FineWebEduDataset(tok, context_length=1024)
loader = DataLoader(dataset, batch_size=8, num_workers=4)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
device = "cuda" if torch.cuda.is_available() else "cpu"

losses = Trainer(
    model,
    optimizer,
    loader,
    device,
    context_length=1024,
    use_amp=device == "cuda",
).train(epochs=1, max_steps=20_000)

import torch

from olm.nn.attention import AttentionBase

class LocalWindowAttention(AttentionBase):
    def __init__(self, embed_dim, num_heads, window=256):
        super().__init__(embed_dim, num_heads)
        self.window = window

    def compute_attention(self, q, k, v, mask=None):
        scores = (q @ k.transpose(-2, -1)) * self.scale
        seq = q.size(-2)
        pos = torch.arange(seq, device=q.device)
        local = (pos[:, None] - pos[None, :]).abs() <= self.window
        causal = pos[:, None] >= pos[None, :]
        scores = scores.masked_fill(~(local & causal), float("-inf"))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        probs = self.dropout(scores.softmax(dim=-1))
        return probs @ v

attention = LocalWindowAttention(d_model, heads, window=256)
# Drop it into a Block, a custom model, or your PyTorch loop.