Spaces:
Sleeping
Sleeping
File size: 5,046 Bytes
59e2c8a 871f869 59e2c8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """
utils/model_loader.py
─────────────────────
Load a local (or HF Hub) model via HuggingFace Transformers.
Returns a model_bundle dict that every benchmark consumes.
"""
from __future__ import annotations
import torch
from typing import Any
DTYPE_MAP = {
"float32": torch.float32,
"float16": torch.float16,
"bfloat16": torch.bfloat16,
}
def load_model(
model_path: str,
device: str = "auto",
dtype: str = "bfloat16",
model_type: str | None = "auto",
) -> dict[str, Any]:
"""
Load model + tokenizer from a local path or HF Hub ID.
Returns
-------
model_bundle : dict with keys
model – the loaded AutoModelForCausalLM
tokenizer – the matching AutoTokenizer
device – resolved torch device string
dtype – resolved torch dtype
param_count – float (billions)
model_path – original path string
generate_fn – convenience callable (prompt → str)
"""
if model_type and model_type not in ("auto", "hf"):
raise ValueError(
f"Unsupported model_type {model_type!r}; use 'auto' or 'hf'."
)
# ── lazy imports so the module is importable without torch installed ──────
try:
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
)
except ImportError as e:
raise ImportError(
"transformers is required: pip install transformers accelerate"
) from e
model_path = str(model_path)
# ── Quantization config ───────────────────────────────────────────────────
quant_cfg = None
torch_dtype = DTYPE_MAP.get(dtype, torch.bfloat16)
if dtype == "int4":
quant_cfg = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
torch_dtype = None
elif dtype == "int8":
quant_cfg = BitsAndBytesConfig(load_in_8bit=True)
torch_dtype = None
# ── Load tokenizer ────────────────────────────────────────────────────────
tokenizer = AutoTokenizer.from_pretrained(
model_path,
trust_remote_code=True,
padding_side="left",
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# ── Load model ────────────────────────────────────────────────────────────
model = AutoModelForCausalLM.from_pretrained(
model_path,
device_map=device,
torch_dtype=torch_dtype,
quantization_config=quant_cfg,
trust_remote_code=True,
)
model.eval()
# ── Parameter count ───────────────────────────────────────────────────────
param_count = sum(p.numel() for p in model.parameters()) / 1e9
# ── Resolve actual device ─────────────────────────────────────────────────
resolved_device = next(model.parameters()).device
# ── Convenience generate function ─────────────────────────────────────────
def generate_fn(
prompt: str,
max_new_tokens: int = 512,
temperature: float = 0.0,
stop_strings: list[str] | None = None,
) -> str:
"""Run inference and return the decoded completion (without prompt)."""
inputs = tokenizer(prompt, return_tensors="pt").to(resolved_device)
gen_kwargs: dict[str, Any] = dict(
**inputs,
max_new_tokens=max_new_tokens,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
if temperature > 0:
gen_kwargs.update(do_sample=True, temperature=temperature, top_p=0.95)
else:
gen_kwargs["do_sample"] = False
with torch.no_grad():
output_ids = model.generate(**gen_kwargs)
# Strip the input tokens from output
new_ids = output_ids[0][inputs["input_ids"].shape[1]:]
return tokenizer.decode(new_ids, skip_special_tokens=True).strip()
return {
"model": model,
"tokenizer": tokenizer,
"device": str(resolved_device),
"dtype": dtype,
"param_count": param_count,
"model_path": model_path,
"model_type": "hf",
"generate_fn": generate_fn,
}
|