| """Model/tokenizer loading and stop-token resolution.""" | |
| from __future__ import annotations | |
| from pathlib import Path | |
| from ..paths import BASE_MODEL | |
| def resolve_model_dir(model: str | Path) -> tuple[str, str | None]: | |
| """Return (model_path, adapter_path). Prefers merged weights when present.""" | |
| p = Path(model) | |
| if p.is_dir(): | |
| if (p / "merged").is_dir() and (p / "merged" / "config.json").is_file(): | |
| return str(p / "merged"), None | |
| if (p / "config.json").is_file(): | |
| return str(p), None | |
| if (p / "adapter_config.json").is_file(): | |
| return BASE_MODEL, str(p) | |
| return str(model), None | |
| def load_model_and_tokenizer(model: str | Path, *, prefer_cpu: bool = False): | |
| """Load HF causal LM. Set prefer_cpu=True to force CPU float32 (no CUDA map).""" | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| model_path, adapter_path = resolve_model_dir(model) | |
| tokenizer = AutoTokenizer.from_pretrained(adapter_path or model_path, trust_remote_code=True) | |
| use_cuda = torch.cuda.is_available() and not prefer_cpu | |
| dtype = torch.bfloat16 if use_cuda else torch.float32 | |
| device_map = "auto" if use_cuda else "cpu" | |
| lm = AutoModelForCausalLM.from_pretrained( | |
| model_path, trust_remote_code=True, device_map=device_map, torch_dtype=dtype | |
| ) | |
| if adapter_path: | |
| from peft import PeftModel | |
| lm = PeftModel.from_pretrained(lm, adapter_path).merge_and_unload() | |
| lm.eval() | |
| return lm, tokenizer | |
| def build_stop_ids(tokenizer) -> list[int]: | |
| """ChatML / LFM end tokens for clean stopping.""" | |
| candidates = ["<|im_end|>", "<|endoftext|>", tokenizer.eos_token] | |
| ids: list[int] = [] | |
| for s in candidates: | |
| if not s: | |
| continue | |
| try: | |
| tid = tokenizer.convert_tokens_to_ids(s) | |
| if tid is not None and tid != tokenizer.unk_token_id and int(tid) >= 0: | |
| ids.append(int(tid)) | |
| except Exception: | |
| pass | |
| try: | |
| enc = tokenizer.encode(s, add_special_tokens=False) | |
| if len(enc) == 1: | |
| ids.append(int(enc[0])) | |
| except Exception: | |
| pass | |
| out: list[int] = [] | |
| for i in ids: | |
| if i not in out: | |
| out.append(i) | |
| if tokenizer.eos_token_id is not None and tokenizer.eos_token_id not in out: | |
| out.append(int(tokenizer.eos_token_id)) | |
| return out | |