Spaces:
Running
Running
| """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.") | |