"""Interactive zh<->en translation demo for Interpreter-Qwen3-1.7B.

    pip install vllm
    python translate.py                      # loads Ismantic/Interpreter-Qwen3-1.7B
    python translate.py --model_path ./       # or a local snapshot dir

Type a sentence + Enter; direction (zh->en / en->zh) is auto-detected by whether the
line contains Chinese characters. Ctrl-C / Ctrl-D to quit. Uses the model's ChatML
prompt format with <|im_end|> as the stop token and greedy decoding.
"""
import sys
import argparse
from vllm import LLM, SamplingParams

PROMPT_ZH2EN = "Translate the following text from Chinese to English.\nChinese: {src}\nEnglish:"
PROMPT_EN2ZH = "Translate the following text from English to Chinese.\nEnglish: {src}\nChinese:"


def is_zh(s):
    return any("一" <= c <= "鿿" for c in s)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model_path", default="Ismantic/Interpreter-Qwen3-1.7B")
    ap.add_argument("--gpu_mem", type=float, default=0.5)
    args = ap.parse_args()

    llm = LLM(model=args.model_path, gpu_memory_utilization=args.gpu_mem, max_model_len=1024)
    sp = SamplingParams(max_tokens=256, temperature=0, stop=["<|im_end|>"])

    print("\n>>> Type a sentence and press Enter (zh<->en auto-detected). Ctrl-C / Ctrl-D to quit.\n")
    try:
        for line in sys.stdin:
            src = line.strip()
            if not src:
                continue
            tmpl = PROMPT_ZH2EN if is_zh(src) else PROMPT_EN2ZH
            prompt = f"<|im_start|>user\n{tmpl.format(src=src)}<|im_end|>\n<|im_start|>assistant\n"
            out = llm.generate([prompt], sp, use_tqdm=False)[0].outputs[0].text.strip()
            print("->", out, "\n")
    except (KeyboardInterrupt, EOFError):
        print("\nbye")


if __name__ == "__main__":
    main()