model_loader.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Tuple

SUPPORTED_MODELS = {
    "gpt2":               "GPT-2 Small (124M) — default, ultra-lightweight",
    "gpt2-medium":        "GPT-2 Medium (345M) — better distributions",
    "Qwen/Qwen2.5-0.5B":  "Qwen2.5 0.5B — modern, Apache 2.0",
    "Qwen/Qwen2.5-1.5B":  "Qwen2.5 1.5B — best quality on CPU",
    "HuggingFaceTB/SmolLM2-135M": "SmolLM2 135M — ultra-portable",
}


def get_device() -> str:
    """Auto-detect best available device."""
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def load_model(
    model_name: str = "gpt2",
    device: str = None,
) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
    """
    Load a causal LM and its tokenizer.

    Args:
        model_name: HuggingFace model ID (default: gpt2)
        device: 'cpu' or 'cuda'. If None, auto-detects.

    Returns:
        (model, tokenizer, device)
    """
    if device is None:
        device = get_device()

    print(f"[ModelLoader] Loading '{model_name}' on {device.upper()}...")

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Ensure pad token exists
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True,
    )
    model.to(device)
    model.eval()

    print(f"[ModelLoader] ✓ Model loaded. Vocab size: {model.config.vocab_size:,}")
    return model, tokenizer, device