"""The model card's example: two steps of the questioner on a two-hop question."""

import re

import torch
from huggingface_hub import hf_hub_download
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "dolev31/ProactiveInquirer-Qwen3-8B"
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, REPO)  # seed 1; add subfolder="seed2" for the second seed

template = open(hf_hub_download(REPO, "prompts/inquirer_prompted.txt"), encoding="utf-8").read()
placebo = open(hf_hub_download(REPO, "prompts/fragment_user_channel_placebo.txt"), encoding="utf-8").read()


def next_action(**state):
    fields = dict(state, user_channel=placebo.strip())
    prompt = re.sub(r"\{\{(\w+)\}\}", lambda m: str(fields[m.group(1)]), template)
    ids = tok.apply_chat_template(
        [{"role": "user", "content": prompt}],
        add_generation_prompt=True,
        enable_thinking=False,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)
    out = model.generate(**ids, max_new_tokens=200, do_sample=False)
    return tok.decode(out[0, ids["input_ids"].shape[1] :], skip_special_tokens=True)


question = "Who was the spouse of the director of the film The Great Flamarion?"
instructions = (
    "Answer the question using a closed pool of 20 paragraphs. You may issue retrieval "
    "queries against that pool before answering; several paragraphs are distractors, and "
    "the answer usually requires composing facts from more than one of them."
)
print(next_action(
    question=question, instructions=instructions, evidence="(nothing retrieved yet)",
    draft="(no draft yet)", history="(nothing asked yet)",
))
evidence = (
    "[3f2a9c1b7d4e] The Great Flamarion\n"
    "The Great Flamarion is a 1945 American film noir directed by Anthony Mann and starring "
    "Erich von Stroheim, Mary Beth Hughes and Dan Duryea."
)
print(next_action(
    question=question, instructions=instructions, evidence=evidence,
    draft="The film was directed by Anthony Mann; his spouse is not yet known.",
    history="Q1: Who directed the film The Great Flamarion?\n"
            "A1: The Great Flamarion (1945) was directed by Anthony Mann.",
))