"""End-to-end coherence check: a few natural-language prompts, printed as decoded text.

`h20_ask.py` takes prompts on the command line, which a non-UTF-8 login shell mangles into
`????`. The prompts therefore live in this file, so the text the model actually sees is the
text written here.

    python scripts/h20_state_check.py --max_tokens 64            # greedy
    python scripts/h20_state_check.py --temperature 0.99         # real sampling (sample path)
"""

import argparse
import os

from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env
from _launch import launch_dp

from deepinfer.sampling_params import SamplingParams

PROMPTS = [
    "问题：中国的首都是哪里？\n答案：",
    "请用两句话解释什么是注意力机制。\n答：",
    "写一句描写秋天的诗。\n",
]

parser = argparse.ArgumentParser()
parser.add_argument("--max_tokens", type=int, default=64)
parser.add_argument("--temperature", type=float, default=0.0)
parser.add_argument("--ep", type=int, default=8)
parser.add_argument("--cp", type=int, default=1)
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME)
args = parser.parse_args()

ep_size, cp_size = args.ep, args.cp
dp_size = max(ep_size // cp_size, 1)
world_size = dp_size * cp_size
assert world_size % ep_size == 0


def main(dp_size: int, dp_rank: int) -> None:
    setup_common_env(
        model_name=args.model_name,
        sampler_force_topk=4096,
        set_override_experts=False,
        enable_eplb=False,
    )
    os.environ.setdefault("DEEPINFER_BATCH_INVARIANT", "1")

    llm = create_llm(
        args.model_name,
        dp_rank,
        ep_size=ep_size,
        cp_size=cp_size,
        env={"DEEPINFER_LOGGING_INTERVAL_STEP": "10"},
        data_parallel_size=dp_size,
        max_num_batched_tokens=4096,
        max_model_len=4096,
        max_num_seqs=max(1, len(PROMPTS)),
        gpu_memory_utilization=0.85,
        lookahead_tokens=0,
    )

    sampling_params = [
        SamplingParams(temperature=args.temperature, max_tokens=args.max_tokens, ignore_eos=True, seq_seed=0) for _ in PROMPTS
    ]
    outputs = llm.generate(prompts=PROMPTS, sampling_params=sampling_params, use_tqdm=False)
    for i, output in enumerate(outputs):
        assert output, f"sample {i}: no generated tokens"
        tokens = [x.token_id for x in output]
        text = llm.llm_engine.tokenizer.decode(tokens)
        print(f"\n=== DP {dp_rank} sample {i}  temp={args.temperature}")
        print(f"=== prompt: {PROMPTS[i]!r}")
        print(f"=== token ids ({len(tokens)}): {tokens}")
        print(f"=== continuation: {text!r}", flush=True)

    llm.close()


if __name__ == "__main__":
    launch_dp(main, dp_size)