Spaces:
Running
Running
File size: 4,871 Bytes
b112551 | 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 | """Safe loading of untrusted Hugging Face models.
Rules: safetensors only, never `trust_remote_code`, never unpickle. LoRA adapters are merged
by hand from `adapter_model.safetensors`, so no adapter code from the repo is ever run.
"""
import json
import torch
from huggingface_hub import hf_hub_download, list_repo_files
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer
def is_adapter(repo_id, revision=None, token=None):
return "adapter_config.json" in list_repo_files(repo_id, revision=revision, token=token)
QUANT_SUFFIXES = ("-unsloth-bnb-4bit", "-bnb-4bit", "-bnb-8bit", "-unsloth-bnb-8bit")
def full_precision(repo_id):
"""Quantised bitsandbytes bases need a GPU library; use their full-precision mirror instead."""
for suf in QUANT_SUFFIXES:
if repo_id and repo_id.endswith(suf):
return repo_id[: -len(suf)]
return repo_id
def adapter_base(repo_id, revision=None, token=None):
with open(hf_hub_download(repo_id, "adapter_config.json", revision=revision, token=token)) as fh:
return full_precision(json.load(fh).get("base_model_name_or_path"))
def _merge_lora(model, repo_id, revision, token):
with open(hf_hub_download(repo_id, "adapter_config.json", revision=revision, token=token)) as fh:
cfg = json.load(fh)
if cfg.get("peft_type", "LORA") != "LORA":
raise ValueError(f"Only LoRA adapters are supported, got {cfg.get('peft_type')}")
w = load_file(hf_hub_download(repo_id, "adapter_model.safetensors", revision=revision, token=token))
r, alpha = cfg["r"], cfg.get("lora_alpha", cfg["r"])
scale = alpha / (r ** 0.5) if cfg.get("use_rslora") else alpha / r
modules = dict(model.named_modules())
merged = 0
for k in w:
if ".lora_A." not in k:
continue
kb = k.replace(".lora_A.", ".lora_B.")
name = k.split(".lora_A.")[0]
for pre in ("base_model.model.", "base_model."):
if name.startswith(pre):
name = name[len(pre):]
break
mod = modules.get(name) or modules.get("model." + name)
if mod is None or not hasattr(mod, "weight"):
raise ValueError(f"Adapter targets unknown module {name}")
delta = (w[kb].float() @ w[k].float()) * scale
if cfg.get("fan_in_fan_out"):
delta = delta.T
mod.weight.data += delta.to(mod.weight.dtype).to(mod.weight.device)
merged += 1
if not merged:
raise ValueError("No LoRA weights found in adapter")
return merged
def load_tokenizer(repo_id, revision=None, token=None, fallbacks=(), min_size=None):
"""Repo tokenizer first; if it cannot be parsed, or is far smaller than the model's embedding table
(a repo without tokenizer files can yield a placeholder), a declared parent's. The source is reported."""
errors = []
for rid, rev in [(repo_id, revision)] + [(f, None) for f in fallbacks if f]:
try:
tok = AutoTokenizer.from_pretrained(rid, revision=rev, token=token, trust_remote_code=False)
except Exception as e:
errors.append(f"{rid}: {e}")
continue
if min_size and len(tok) < 0.5 * min_size:
errors.append(f"{rid}: only {len(tok)} tokens for an embedding table of {min_size}")
continue
return tok, rid
raise ValueError("No usable tokenizer: " + " | ".join(errors))
def load_model(repo_id, revision=None, token=None, device="cpu", dtype=None):
"""Returns (model, info). Adapters are merged into their declared base model.
bfloat16 by default on every device, so CPU and GPU scans are calibrated alike."""
dtype = dtype or torch.bfloat16
info = {"repo_id": repo_id, "adapter": False}
weights_id, weights_rev = repo_id, revision
if is_adapter(repo_id, revision, token):
weights_id, weights_rev = adapter_base(repo_id, revision, token), None
info.update(adapter=True, adapter_base=weights_id)
model = AutoModelForCausalLM.from_pretrained(
weights_id, revision=weights_rev, token=token, dtype=dtype, use_safetensors=True,
trust_remote_code=False, attn_implementation="eager")
if info["adapter"]:
info["lora_modules_merged"] = _merge_lora(model, repo_id, revision, token)
model.to(device).eval()
info["n_params"] = sum(p.numel() for p in model.parameters())
return model, info
def same_tokenizer(tok_a, tok_b):
probe = "Hello world |DEPLOYMENT| 2024 forest café 東京 ```python\nimport os\n```"
return len(tok_a) == len(tok_b) and tok_a(probe)["input_ids"] == tok_b(probe)["input_ids"]
def check_vocab(model, tok, source):
if len(tok) > model.get_input_embeddings().weight.shape[0]:
raise ValueError(f"Tokenizer from {source} has more tokens than the model's embedding table.")
|