mech_workbench/model_loader.py

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

def load_model_and_tokenizer(model_name: str, device: str = "auto", dtype: str = "auto"):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    torch_dtype = "auto" if dtype == "auto" else getattr(torch, dtype)

    model_kwargs: dict[str, Any] = {"torch_dtype": torch_dtype}
    if device == "auto":
        model_kwargs["device_map"] = "auto"

    model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
    if device != "auto":
        model.to(device)
    model.eval()
    return model, tokenizer

def get_transformer_layers(model: Any) -> Sequence[Any]:
    candidates = [
        ("model", "layers"),
        ("transformer", "h"),
        ("gpt_neox", "layers"),
        ("model", "decoder", "layers"),
    ]

    for path in candidates:
        current = model
        try:
            for attr in path:
                current = getattr(current, attr)
        except AttributeError:
            continue
        if hasattr(current, "__len__") and hasattr(current, "__getitem__"):
            return current

    raise AttributeError(
        "could not locate transformer block list; add this architecture to "
        "get_transformer_layers"
    )