| |
| """vllm predict with extras (vllm 0.19 compatible).""" |
| import os, json, argparse, time, pathlib |
| os.environ.setdefault("HF_HOME", "/mnt/msrh/Magic_submission/hf_cache") |
| from vllm import LLM, SamplingParams |
| from vllm.lora.request import LoRARequest |
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base", required=True) |
| ap.add_argument("--adapter", required=True) |
| ap.add_argument("--rag_test", default="/mnt/msrh/Magic_submission/LF/data/msrh_rag_test_k3_AfriE5_TV.json") |
| ap.add_argument("--out_jsonl", required=True) |
| ap.add_argument("--max_lora_rank", type=int, default=128) |
| ap.add_argument("--max_new", type=int, default=512) |
| ap.add_argument("--max_model_len", type=int, default=2560) |
| ap.add_argument("--max_num_seqs", type=int, default=256) |
| ap.add_argument("--mem_util", type=float, default=0.85) |
| ap.add_argument("--temperature", type=float, default=0.0) |
| ap.add_argument("--top_p", type=float, default=1.0) |
| ap.add_argument("--repetition_penalty", type=float, default=1.0) |
| ap.add_argument("--frequency_penalty", type=float, default=0.0) |
| ap.add_argument("--best_of", type=int, default=1, help="n samples; pick highest cumulative_logprob") |
| ap.add_argument("--no_think", action="store_true", help="disable thinking mode (e.g., Qwen3.5)") |
| ap.add_argument("--use_beam", action="store_true") |
| ap.add_argument("--beam_width", type=int, default=4) |
| args = ap.parse_args() |
|
|
| print("loading", flush=True) |
| llm = LLM(model=args.base, enable_lora=True, max_loras=1, max_lora_rank=args.max_lora_rank, |
| gpu_memory_utilization=args.mem_util, max_model_len=args.max_model_len, max_num_seqs=args.max_num_seqs, |
| dtype="bfloat16", trust_remote_code=True) |
| tok = llm.get_tokenizer() |
| rows = [json.loads(l) for l in open(args.rag_test)] |
| prompts = [tok.apply_chat_template(r["messages"][:1], tokenize=False, add_generation_prompt=True, enable_thinking=not args.no_think) for r in rows] |
| print(f"prompts: {len(prompts)}", flush=True) |
| lr = LoRARequest("rag", 1, args.adapter) |
| t0 = time.time() |
|
|
| if args.use_beam: |
| from vllm.sampling_params import BeamSearchParams |
| bp = BeamSearchParams(beam_width=args.beam_width, max_tokens=args.max_new, temperature=0.0) |
| try: |
| outs = llm.beam_search(prompts, bp, lora_request=lr) |
| except TypeError: |
| outs = llm.beam_search(prompts, bp) |
| out_texts = [] |
| for o in outs: |
| try: |
| seq = o.sequences[0] |
| txt = getattr(seq, "text", None) or tok.decode(seq.tokens, skip_special_tokens=True) |
| except Exception: |
| txt = "" |
| out_texts.append(txt) |
| else: |
| sp = SamplingParams( |
| n=args.best_of, |
| temperature=args.temperature, top_p=args.top_p, |
| max_tokens=args.max_new, repetition_penalty=args.repetition_penalty, frequency_penalty=args.frequency_penalty, |
| ) |
| outs = llm.generate(prompts, sp, lora_request=lr) |
| out_texts = [] |
| for o in outs: |
| cands = o.outputs |
| if len(cands) == 1: |
| out_texts.append(cands[0].text) |
| else: |
| |
| best = max(cands, key=lambda c: c.cumulative_logprob if c.cumulative_logprob is not None else -1e18) |
| out_texts.append(best.text) |
|
|
| print(f"done {(time.time()-t0)/60:.1f}min", flush=True) |
| p = pathlib.Path(args.out_jsonl); p.parent.mkdir(parents=True, exist_ok=True) |
| with open(p, "w") as f: |
| for prompt, txt in zip(prompts, out_texts): |
| f.write(json.dumps({"prompt": prompt, "predict": txt}, ensure_ascii=False) + "\n") |
| print("wrote") |
|
|
