Aurora

Model was pre-trained for 20 epochs on enwik8 and achieved a validation loss of 1.3597 1.0565, a train loss of 1.3171, and a BPB of 1.5242.

Model is 23M parameters.

Inference Code:

import torch
from pathlib import Path
from x_transformers import TransformerWrapper, Decoder, AutoregressiveWrapper


class InferenceConfig:
    num_tokens: int = 256
    max_seq_len: int = 1024
    model_dim: int = 512
    model_depth: int = 6
    model_heads: int = 6
    model_path: str = "best_model.pth"


def count_parameters(model: torch.nn.Module) -> tuple[int, int]:
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total_params, trainable_params


def format_parameters(num_params: int) -> str:
    if num_params >= 1e9:
        return f"{num_params / 1e9:.2f}B"
    elif num_params >= 1e6:
        return f"{num_params / 1e6:.2f}M"
    elif num_params >= 1e3:
        return f"{num_params / 1e3:.2f}K"
    else:
        return f"{num_params}"


def load_model(config: InferenceConfig) -> AutoregressiveWrapper:
    """Load the trained model for inference."""
    model = TransformerWrapper(
        num_tokens=config.num_tokens,
        max_seq_len=config.max_seq_len,
        attn_layers=Decoder(
            dim=config.model_dim,
            depth=config.model_depth,
            heads=config.model_heads,
            use_simple_rmsnorm=True,
            ff_glu=True,
            ff_no_bias=True,
            attn_qk_norm=True,
            attn_qk_norm_dim_scale=True,
            attn_flash=True,
            attn_kv_heads=2,
        ),
    )

    model = AutoregressiveWrapper(model, mask_prob=0.0).cuda()
    model = model.to(dtype=torch.bfloat16)

    total_params, trainable_params = count_parameters(model)

    checkpoint = torch.load(config.model_path)
    model.load_state_dict(checkpoint["model_state_dict"])

    model.eval()

    print(f"Model loaded from {config.model_path}")
    print(f"Training step: {checkpoint['step']}, Loss: {checkpoint['loss']:.4f}")
    print(f"Total parameters: {format_parameters(total_params)} ({total_params:,})")
    print(
        f"Trainable parameters: {format_parameters(trainable_params)} ({trainable_params:,})"
    )

    return model


def generate_sample(
    model: AutoregressiveWrapper,
    config: InferenceConfig,
    prompt: str = "The",
    temperature: float = 1.0,
    top_p: float = 0.9,
    max_length: int = 1024,
) -> str:
    model.eval()

    with torch.no_grad():
        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            prompt_bytes = prompt.encode("utf-8")
            prompt_tokens = torch.tensor(
                [[b for b in prompt_bytes]], dtype=torch.long
            ).cuda()

            generation_length = max_length - len(prompt_bytes)

            generated = model.generate(
                prompt_tokens,
                seq_len=generation_length,
                temperature=temperature,
                filter_logits_fn="top_p",
                filter_kwargs={"thres": top_p},
                cache_kv=True,
            )

    generated_bytes = bytes(generated[0].cpu().numpy())
    try:
        generated_text = generated_bytes.decode("utf-8", errors="replace")
    except:
        generated_text = str(generated_bytes)

    return generated_text


def main():
    config = InferenceConfig()

    # Check if model exists
    if not Path(config.model_path).exists():
        print(f"Error: Model file '{config.model_path}' not found!")
        return

    # Load model
    print("Loading model...")
    model = load_model(config)

    print("=" * 80)

    sample = generate_sample(
        model,
        config,
        prompt="Here be dragons",
        temperature=1.0,
        top_p=0.9,
        max_length=1024,
    )

    print(sample)
    print("=" * 80)


if __name__ == "__main__":
    main()