#!/usr/bin/env python3
"""
多句并发推理演示

演示如何使用 Transformers 库的 batch 推理功能，
在单次前向传播中处理多条输入，观察性能提升。
"""

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time

# ============================================================================
# 模型加载
# ============================================================================
print("Loading model...")
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 关键：Decoder-only 模型（如 Qwen、GPT）必须使用左填充而非右填充
# 原因：自回归解码模型只能看到当前位置之前的 token
# - 右填充：[有效输入] + [padding] → 模型会错误地尝试生成 padding token
# - 左填充：[padding] + [有效输入] → 模型可以安全地忽略前面的 padding
# 正确设置 padding_side 可以避免警告，同时保证生成质量
tokenizer.padding_side = "left"

model = AutoModelForCausalLM.from_pretrained(model_name, dtype=torch.bfloat16)
model.eval()
print("Model loaded!\n")

# ============================================================================
# 定义三条输入 prompts
# ============================================================================
# 三条不同的用户问题，展示模型如何处理多样化的输入
messages_list = [
    # 第一条：数学问题
    [{"role": "user", "content": "请计算：1 加 3 等于几？"}],
    
    # 第二条：创意写作
    [{"role": "user", "content": "写一句关于春天的诗"}],
    
    # 第三条：常识问题
    [{"role": "user", "content": "中国的首都是哪个城市？"}],
]

print("=" * 70)
print("📋 输入的三条 Prompts")
print("=" * 70)
for idx, messages in enumerate(messages_list, 1):
    print(f"\n[Prompt {idx}] {messages[0]['content']}")

# ============================================================================
# 方式 1：单条推理（对比基线）
# ============================================================================
print("\n\n" + "=" * 70)
print("⏱️  方式 1：单条推理（顺序执行）")
print("=" * 70)

start_time = time.time()

responses_sequential = []
for idx, messages in enumerate(messages_list, 1):
    # 格式化输入
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    # 单条推理
    model_inputs = tokenizer([text], return_tensors="pt")
    with torch.no_grad():
        generated_ids = model.generate(
            model_inputs["input_ids"],
            attention_mask=model_inputs.get("attention_mask"),
            max_new_tokens=64,
            temperature=0.7,
            top_p=0.95,
            do_sample=True
        )
    
    response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    response = response.split("assistant\n")[-1] if "assistant" in response else response
    responses_sequential.append(response)
    
    print(f"\n[Output {idx}] {response}")

sequential_time = time.time() - start_time
print(f"\n⏱️  总耗时：{sequential_time:.2f} 秒")

# ============================================================================
# 方式 2：Batch 推理（并发处理）
# ============================================================================
print("\n\n" + "=" * 70)
print("⚡ 方式 2：Batch 推理（并发处理）")
print("=" * 70)

start_time = time.time()

# 步骤 1：格式化所有输入
formatted_prompts = []
for messages in messages_list:
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    formatted_prompts.append(text)

# 步骤 2：将所有输入 tokenize 成 batch
# tokenizer 会自动添加 padding 使所有序列长度一致
model_inputs = tokenizer(
    formatted_prompts,
    return_tensors="pt",
    padding=True,  # 启用 padding，短的序列会被填充到最长的长度
    truncation=True,  # 长序列会被截断
)

# 步骤 3：一次性推理整个 batch
with torch.no_grad():
    generated_ids = model.generate(
        model_inputs["input_ids"],
        attention_mask=model_inputs.get("attention_mask"),
        max_new_tokens=64,
        temperature=0.7,
        top_p=0.95,
        do_sample=True
    )

# 步骤 4：解码所有生成结果
responses_batch = []
for idx, generated_id in enumerate(generated_ids):
    response = tokenizer.decode(generated_id, skip_special_tokens=True)
    response = response.split("assistant\n")[-1] if "assistant" in response else response
    responses_batch.append(response)
    
    print(f"\n[Output {idx + 1}] {response}")

batch_time = time.time() - start_time
print(f"\n⏱️  总耗时：{batch_time:.2f} 秒")

# ============================================================================
# 性能对比分析
# ============================================================================
print("\n\n" + "=" * 70)
print("📊 性能对比分析")
print("=" * 70)

speedup = sequential_time / batch_time
print(f"\n单条推理耗时：   {sequential_time:.2f} 秒")
print(f"Batch 推理耗时： {batch_time:.2f} 秒")
print(f"性能提升倍数：   {speedup:.2f}x")
print(f"总体加速：       {(1 - batch_time / sequential_time) * 100:.1f}%")