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" )