#!/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")