| |
| """Load an EviSuff adapter stack and optionally run a smoke-test prompt.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
|
|
| ADAPTERS = ("answer-sft", "no-gate", "full-evisuff") |
| BASE_MODEL = "Qwen/Qwen3-8B" |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--adapter", choices=ADAPTERS, default="full-evisuff") |
| parser.add_argument( |
| "--repo-id", |
| help="Hugging Face model repository. Omit to use the local repository clone.", |
| ) |
| parser.add_argument("--base-model", default=BASE_MODEL) |
| parser.add_argument("--revision", default=None, help="Optional base-model revision.") |
| parser.add_argument("--prompt", help="Optional prompt for a short generation smoke test.") |
| parser.add_argument("--max-new-tokens", type=int, default=128) |
| return parser.parse_args() |
|
|
|
|
| def adapter_location(repo_id: str | None, adapter: str) -> str: |
| if repo_id: |
| return f"{repo_id}/{adapter}" |
| return str(Path(__file__).resolve().parents[1] / adapter) |
|
|
|
|
| def load_stack(args: argparse.Namespace): |
| import torch |
| from peft import PeftModel |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 |
| base = AutoModelForCausalLM.from_pretrained( |
| args.base_model, |
| revision=args.revision, |
| torch_dtype=dtype, |
| device_map="auto", |
| ) |
| tokenizer = AutoTokenizer.from_pretrained(args.base_model, revision=args.revision) |
|
|
| if args.repo_id: |
| answer_model = PeftModel.from_pretrained( |
| base, |
| args.repo_id, |
| subfolder="answer-sft", |
| ) |
| else: |
| answer_model = PeftModel.from_pretrained( |
| base, |
| adapter_location(None, "answer-sft"), |
| ) |
|
|
| if args.adapter == "answer-sft": |
| return answer_model, tokenizer |
|
|
| merged_answer = answer_model.merge_and_unload() |
| if args.repo_id: |
| model = PeftModel.from_pretrained( |
| merged_answer, |
| args.repo_id, |
| subfolder=args.adapter, |
| ) |
| else: |
| model = PeftModel.from_pretrained( |
| merged_answer, |
| adapter_location(None, args.adapter), |
| ) |
| return model, tokenizer |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| model, tokenizer = load_stack(args) |
| print(f"Loaded {args.adapter} with the required adapter stack.") |
| if not args.prompt: |
| return |
|
|
| inputs = tokenizer(args.prompt, return_tensors="pt").to(model.device) |
| output = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False) |
| print(tokenizer.decode(output[0], skip_special_tokens=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
