在各种 AI里反复切换，科研真的更高效了吗？

"""
简化版 Speculative Decoding 示例

本示例模拟了 speculative decoding 的核心交互逻辑：
- DraftModel（草稿模型）快速生成候选 token。
- MainModel（主模型）验证这些候选 token 是否可接受。
- 每轮只生成一个 token（k=1），便于理解基本流程。
"""

import random


class MainModel:
    """
    主模型（Target Model）：负责高质量地验证 token。
    在真实场景中，该模型通常较大、较慢，但生成质量高。
    """

    def __init__(self, seed=42):
        # 设置随机种子以保证可复现性（仅用于演示）
        random.seed(seed)

    def validate_token(self, token, context):
        """
        验证草稿模型生成的 token 是否可被主模型接受。

        为简化，本示例采用概率性接受策略：
        - 如果 token 与主模型“期望”的 token 一致，则接受；
        - 否则以一定概率拒绝（模拟分布差异）。

        在真实 speculative decoding 中，此处会执行拒绝采样逻辑。

        参数:
            token (int): 待验证的候选 token ID。
            context (list): 当前已生成的 token 序列（上下文）。

        返回:
            bool: True 表示接受，False 表示拒绝。
        """
        # 模拟主模型对下一个 token 的“偏好”：
        # 为简单起见，我们假设主模型期望的 token 是 (sum(context) + 1) % 10
        expected_token = (sum(context) + 1) % 10

        # 如果草稿 token 与期望一致，直接接受
        if token == expected_token:
            return True

        # 否则，以较低概率接受（模拟拒绝采样中的“小概率接受”）
        # 这里仅为演示，实际中会基于 p(x)/q(x) 计算接受概率
        accept_prob = 0.2  # 20% 概率接受错误 token
        return random.random() < accept_prob


class DraftModel:
    """
    草稿模型（Draft Model）：快速生成候选 token。
    在真实场景中，该模型通常较小、较快，但质量略低。
    """

    def __init__(self, seed=123):
        random.seed(seed)

    def generate_tokens(self, context, num_tokens=1):
        """
        基于当前上下文生成若干候选 token。

        为简化，本示例每次只生成一个 token（num_tokens=1），
        且生成逻辑为：随机选择一个接近主模型期望的值。

        参数:
            context (list): 当前已生成的 token 序列。
            num_tokens (int): 要生成的 token 数量（本示例固定为 1）。

        返回:
            list: 生成的候选 token 列表（长度为 num_tokens）。
        """
        # 草稿模型“猜测”主模型可能喜欢的 token
        # 简单策略：基于上下文和一点噪声生成
        base = (sum(context) + 1) % 10 if context else 0
        # 加入 ±1 的扰动，模拟草稿模型的不完美
        candidate = (base + random.choice([-1, 0, 1])) % 10
        return [candidate]


def main():
    """
    主程序：演示 speculative decoding 的三轮交互过程。
    """
    # 初始化模型
    main_model = MainModel(seed=42)
    draft_model = DraftModel(seed=123)

    # 初始上下文（可视为 prompt）
    context = [0]  # 假设初始 token 为 0

    print("初始上下文:", context)
    print("-" * 50)

    # 执行三轮 speculative decoding
    for round_num in range(1, 4):
        print(f"第 {round_num} 轮:")

        # 步骤 1: 草稿模型生成一个候选 token
        candidate_tokens = draft_model.generate_tokens(context, num_tokens=1)
        candidate = candidate_tokens[0]
        print(f"  草稿模型生成候选 token: {candidate}")

        # 步骤 2: 主模型验证该候选 token
        is_accepted = main_model.validate_token(candidate, context)
        print(f"  主模型验证结果: {'接受' if is_accepted else '拒绝'}")

        # 步骤 3: 根据验证结果更新上下文
        if is_accepted:
            context.append(candidate)
            print(f"  接受 token，当前上下文更新为: {context}")
        else:
            # 在真实 speculative decoding 中，此时会从修正分布中重新采样一个 token
            # 为简化，这里我们直接使用主模型“期望”的 token
            fallback_token = (sum(context) + 1) % 10
            context.append(fallback_token)
            print(f"  拒绝候选，回退到主模型推荐 token: {fallback_token}")
            print(f"  当前上下文更新为: {context}")

        print("-" * 50)


if __name__ == "__main__":
    main()

"""
基于 Hugging Face Transformers 的 Speculative Decoding 简化实现

说明：
- MainModel: google/gemma-2b-it（目标模型，高质量）
- DraftModel: TinyLlama/TinyLlama-1.1B-Chat-v1.0（草稿模型，快速）
- 每轮草稿模型生成 1 个候选 token，主模型验证是否接受
- 使用 logits 计算接受概率（简化版拒绝采样）
- 仅用于教学演示，未实现完整 speculative decoding（如多 token 并行验证）
"""

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import random

# 设置设备（自动选择 GPU 或 CPU）
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"使用设备: {device}")

# ==============================
# 1. 加载主模型（Target Model）
# ==============================
print("正在加载主模型 (Gemma-2b) ...")
main_model_name = "google/gemma-2b-it"
main_tokenizer = AutoTokenizer.from_pretrained(main_model_name, trust_remote_code=True)
main_model = AutoModelForCausalLM.from_pretrained(
    main_model_name,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(device).eval()

# ==============================
# 2. 加载草稿模型（Draft Model）
# ==============================
print("正在加载草稿模型 (TinyLlama) ...")
draft_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
draft_tokenizer = AutoTokenizer.from_pretrained(draft_model_name, trust_remote_code=True)
# 确保草稿模型 tokenizer 与主模型兼容（此处仅用于生成候选，不严格对齐）
draft_model = AutoModelForCausalLM.from_pretrained(
    draft_model_name,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(device).eval()

# 添加 pad_token（如果缺失）
if draft_tokenizer.pad_token is None:
    draft_tokenizer.pad_token = draft_tokenizer.eos_token
if main_tokenizer.pad_token is None:
    main_tokenizer.pad_token = main_tokenizer.eos_token

# ==============================
# 3. 辅助函数：从 logits 采样 token
# ==============================
def sample_token(logits, temperature=1.0):
    """从 logits 中采样一个 token ID"""
    probs = torch.softmax(logits / temperature, dim=-1)
    token_id = torch.multinomial(probs, num_samples=1).item()
    return token_id

# ==============================
# 4. Speculative Decoding 主循环
# ==============================
def speculative_decode(
    prompt: str,
    main_model, main_tokenizer,
    draft_model, draft_tokenizer,
    max_new_tokens: int = 10,
    draft_k: int = 1,  # 每轮草稿生成 token 数（简化为 1）
    temperature: float = 1.0
):
    """
    简化版 speculative decoding 实现（k=1）
    """
    # 编码初始 prompt
    input_ids = main_tokenizer.encode(prompt, return_tensors="pt").to(device)
    generated_ids = input_ids.clone()

    print(f"\n初始提示: {prompt}\n")
    print("生成过程:")

    for step in range(max_new_tokens):
        # --- 草稿阶段：用草稿模型生成 1 个候选 token ---
        with torch.no_grad():
            draft_outputs = draft_model(generated_ids)
            draft_logits = draft_outputs.logits[:, -1, :]  # 最后一个位置的 logits
            candidate_token = sample_token(draft_logits, temperature)
            candidate_tensor = torch.tensor([[candidate_token]], device=device)
            speculative_ids = torch.cat([generated_ids, candidate_tensor], dim=1)

        # --- 验证阶段：主模型并行计算 logits（当前 + 候选）---
        with torch.no_grad():
            main_outputs = main_model(speculative_ids)
            main_logits = main_outputs.logits[:, -2:, :]  # 倒数第2（当前）和最后（候选）位置

            # 主模型对“当前应生成位置”的 logits（即倒数第2个）
            main_logits_current = main_logits[:, 0, :]
            # 主模型对“候选 token 位置”的 logits（即最后一个）
            main_logits_candidate = main_logits[:, 1, :]

        # --- 拒绝采样：计算接受概率 ---
        # 获取主模型和草稿模型在候选 token 上的概率
        p = torch.softmax(main_logits_candidate, dim=-1)[0, candidate_token].item()
        q = torch.softmax(draft_logits, dim=-1)[0, candidate_token].item()

        # 避免除零
        if q == 0:
            accept_prob = 0.0
        else:
            accept_prob = min(1.0, p / q)

        accepted = random.random() < accept_prob

        if accepted:
            # 接受候选 token
            generated_ids = speculative_ids
            token_str = main_tokenizer.decode([candidate_token], skip_special_tokens=True)
            print(f"[接受] {token_str}")
        else:
            # 拒绝候选，从主模型重新采样
            fallback_token = sample_token(main_logits_current, temperature)
            generated_ids = torch.cat([generated_ids, torch.tensor([[fallback_token]], device=device)], dim=1)
            token_str = main_tokenizer.decode([fallback_token], skip_special_tokens=True)
            print(f"[拒绝 → 重采样] {token_str}")

        # 检查是否生成结束符
        if generated_ids[0, -1].item() == main_tokenizer.eos_token_id:
            print("遇到 EOS，提前结束。")
            break

    # 返回最终生成文本
    output_text = main_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return output_text

# ==============================
# 5. 主程序
# ==============================
if __name__ == "__main__":
    prompt = "Once upon a time,"
    
    try:
        result = speculative_decode(
            prompt=prompt,
            main_model=main_model,
            main_tokenizer=main_tokenizer,
            draft_model=draft_model,
            draft_tokenizer=draft_tokenizer,
            max_new_tokens=15,
            draft_k=1,
            temperature=0.7
        )
        print("\n" + "="*50)
        print("最终生成结果:")
        print(result)
    except Exception as e:
        print(f"运行出错: {e}")
        print("可能原因：模型下载失败、显存不足、或网络问题。")
        print("建议：使用更小模型（如 'gpt2' 和 'distilgpt2'）进行本地测试。")