| |
| import argparse |
| import gc |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast |
|
|
|
|
| def add_path(path): |
| path = str(Path(path).resolve()) |
| if path not in sys.path: |
| sys.path.insert(0, path) |
|
|
|
|
| def link_raven(model_dir, raven_dir): |
| raven_pkg = raven_dir / "raven" |
| if not raven_pkg.is_dir(): |
| raise FileNotFoundError(f"missing {raven_pkg}") |
|
|
| for target_dir in [model_dir, *Path.home().glob(".cache/huggingface/modules/transformers_modules/Quasar-Preview*")]: |
| if not (target_dir / "modeling_quasar_long.py").exists() and target_dir != model_dir: |
| continue |
| link = target_dir / "raven" |
| if link.is_symlink() or not link.exists(): |
| link.unlink(missing_ok=True) |
| link.symlink_to(raven_pkg, target_is_directory=True) |
|
|
|
|
| def load_tokenizer(model_dir): |
| |
| |
| return PreTrainedTokenizerFast( |
| tokenizer_file=str(model_dir / "tokenizer.json"), |
| bos_token="<|startoftext|>", |
| eos_token="<|endoftext|>", |
| pad_token="<|endoftext|>", |
| cls_token="[CLS]", |
| ) |
|
|
|
|
| def set_param(root, name, tensor): |
| obj = root |
| parts = name.split(".") |
| for part in parts[:-1]: |
| obj = obj[int(part)] if part.isdigit() else getattr(obj, part) |
| old = getattr(obj, parts[-1]) |
| setattr(obj, parts[-1], torch.nn.Parameter(tensor, requires_grad=getattr(old, "requires_grad", False))) |
|
|
|
|
| def fill_meta_params(model, model_dir): |
| index = json.loads((model_dir / "model.safetensors.index.json").read_text())["weight_map"] |
| names = [n for n, p in model.named_parameters() if getattr(p, "is_meta", False)] |
| by_file = {} |
| for name in names: |
| by_file.setdefault(index[name], []).append(name) |
|
|
| for filename, keys in sorted(by_file.items()): |
| tensors = load_file(str(model_dir / filename), device="cpu") |
| for key in keys: |
| set_param(model, key, tensors[key].to(torch.bfloat16)) |
| del tensors |
| gc.collect() |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--model-dir", default="/workspace/quasar_preview_run/Quasar-Preview") |
| p.add_argument("--raven-dir", default="/workspace/quasar_preview_run/goombalab-raven") |
| p.add_argument("--prompt", default="Mo Salah is") |
| p.add_argument("--max-new-tokens", type=int, default=24) |
| p.add_argument("--device", default="cuda:0") |
| args = p.parse_args() |
|
|
| model_dir = Path(args.model_dir).resolve() |
| raven_dir = Path(args.raven_dir).resolve() |
| add_path(model_dir) |
| add_path(raven_dir) |
| link_raven(model_dir, raven_dir) |
|
|
| tokenizer = load_tokenizer(model_dir) |
| model = AutoModelForCausalLM.from_pretrained( |
| str(model_dir), |
| trust_remote_code=True, |
| torch_dtype=torch.bfloat16, |
| low_cpu_mem_usage=False, |
| ).eval() |
| fill_meta_params(model, model_dir) |
| model.to(args.device) |
|
|
| enc = tokenizer(args.prompt, return_tensors="pt") |
| enc.pop("token_type_ids", None) |
| enc = {k: v.to(args.device) for k, v in enc.items()} |
|
|
| with torch.inference_mode(): |
| output_ids = model.generate( |
| **enc, |
| max_new_tokens=args.max_new_tokens, |
| do_sample=False, |
| pad_token_id=tokenizer.eos_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| ) |
|
|
| print(tokenizer.decode(output_ids[0], skip_special_tokens=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
