"""本地流式生成示例。首次下载模型；本轮未执行，不报告速度或回答质量。

主线程消费可显示文本，工作线程生成 Token。队列超时不代表模型正常结束。
取消按生成步检查，不能强行终止一个已经在运行的底层矩阵运算。
"""
import queue
import threading
import time
import torch
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, StoppingCriteria,
    StoppingCriteriaList, TextIteratorStreamer,
)

model_id = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.float32
).eval()
messages = [
    {"role": "system", "content": "根据提供的资料回答，缺少资料时说明无法判断。"},
    {"role": "user", "content": "订单 A100 尚未核实质量问题，现在能确认谁承担运费吗？"},
]
inputs = tokenizer.apply_chat_template(
    messages, tokenize=True, add_generation_prompt=True,
    return_dict=True, return_tensors="pt",
)
# 2 秒是单次读取队列的等待上限，用于定期检查线程，不是整个任务超时。
streamer = TextIteratorStreamer(
    tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=2.0,
)
cancel = threading.Event()
result = {}

class Cancelled(StoppingCriteria):
    def __call__(self, input_ids, scores, **kwargs):
        # 每条序列返回一个布尔值；本示例只有一条，无批量取消语义。
        return torch.full((input_ids.shape[0],), cancel.is_set(),
                          dtype=torch.bool, device=input_ids.device)

def produce():
    try:
        # 推理模式是线程局部设置，所以必须放在生成线程内部。
        with torch.inference_mode():
            result["ids"] = model.generate(
                **inputs, max_new_tokens=96, do_sample=False, streamer=streamer,
                stopping_criteria=StoppingCriteriaList([Cancelled()]),
            )
    except Exception as exc:
        result["error"] = exc
        # 异常时 generate 可能来不及关闭流，显式唤醒消费者。
        streamer.on_finalized_text("", stream_end=True)

worker = threading.Thread(target=produce, daemon=True)
worker.start()
started = time.monotonic()
interrupted = False
try:
    while True:
        # 总超时与一次队列等待分开。CPU 很慢时可自行调整此教学预算。
        if time.monotonic() - started > 120:
            interrupted = True
            cancel.set()
            print("\n超过总等待预算，已请求停止生成。")
            break
        try:
            text = next(streamer)
            # 片段可能为空或包含多个 Token，不用事件数当 Token 数。
            print(text, end="", flush=True)
        except queue.Empty:
            if "error" in result:
                break
            continue
        except StopIteration:
            break
except KeyboardInterrupt:
    interrupted = True
    cancel.set()
    print("\n用户取消，当前文本标记为不完整。")
finally:
    # 有限等待不等于强制终止底层运算；守护线程仅适用于这个独立示例。
    worker.join(timeout=3)

print()
if "error" in result:
    raise RuntimeError("生成失败，已显示的文字不是完整结果") from result["error"]
if interrupted or worker.is_alive() or "ids" not in result:
    print("状态：中断或仍在停止中，不提交为最终业务结果。")
else:
    eos = model.generation_config.eos_token_id
    stop_ids = set(eos if isinstance(eos, list) else [eos])
    last_id = result["ids"][0, -1].item()
    print("状态：", "正常结束标记" if last_id in stop_ids else "达到输出上限，可能截断")