File size: 3,811 Bytes
028081c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | #!/usr/bin/env python3
"""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:
# pick highest cumulative_logprob
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")
|