Attention

import numpy as np

# --- 1. 定义模拟参数和固定“训练好”的权重 ---
# 词汇表：包括特殊标记，用于序列开始和结束
vocab = {"<PAD>": 0, "我": 1, "爱": 2, "北京": 3, "你": 6, "好": 7, "世界": 8, "<SOS>": 4, "<EOS>": 5}
id_to_token = {v: k for k, v in vocab.items()}

# 模型维度设置
d_model = 8       # Embedding 维度
d_k = 4           # Query 和 Key 向量维度
d_v = 4           # Value 向量维度

# 模拟训练好的固定 Embedding 权重
# 每一行是一个 Token 的 Embedding 向量
# （实际中这些值是模型学习到的，这里为演示手动设置一些可区分的数值）
fixed_embedding_weights = np.array([
    [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8], # <PAD> (0)
    [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2], # 我 (1)
    [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], # 爱 (2)
    [0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1], # 北京 (3)
    [0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75], # <SOS> (4)
    [0.95, 0.85, 0.75, 0.65, 0.55, 0.45, 0.35, 0.25], # <EOS> (5)
    [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1], # 你 (6)
    [0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2], # 好 (7)
    [0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3], # 世界 (8)
])

# 模拟训练好的固定 Q, K, V 权重矩阵
# 这些矩阵将 Embedding 转换为 Q, K, V
fixed_W_Q = np.array([
    [0.1, 0.2, 0.0, 0.1], [0.0, 0.1, 0.2, 0.0],
    [0.2, 0.0, 0.1, 0.2], [0.1, 0.2, 0.0, 0.1],
    [0.0, 0.1, 0.2, 0.0], [0.2, 0.0, 0.1, 0.2],
    [0.1, 0.2, 0.0, 0.1], [0.0, 0.1, 0.2, 0.0],
])
fixed_W_K = np.array([
    [0.0, 0.1, 0.2, 0.0], [0.1, 0.0, 0.1, 0.2],
    [0.2, 0.1, 0.0, 0.1], [0.0, 0.1, 0.2, 0.0],
    [0.1, 0.0, 0.1, 0.2], [0.2, 0.1, 0.0, 0.1],
    [0.0, 0.1, 0.2, 0.0], [0.1, 0.0, 0.1, 0.2],
])
fixed_W_V = np.array([
    [0.2, 0.1, 0.0, 0.2], [0.1, 0.2, 0.1, 0.0],
    [0.0, 0.1, 0.2, 0.1], [0.2, 0.1, 0.0, 0.2],
    [0.1, 0.2, 0.1, 0.0], [0.0, 0.1, 0.2, 0.1],
    [0.2, 0.1, 0.0, 0.2], [0.1, 0.2, 0.1, 0.0],
])

# --- 2. 核心注意力函数 ---
def calculate_attention(Q_current, K_cache, V_cache):
    """
    计算缩放点积注意力。
    Q_current: 当前步的 Query (形状 1 x d_k)
    K_cache: 累积的 Key 缓存 (形状 L_current x d_k)
    V_cache: 累积的 Value 缓存 (形状 L_current x d_v)
    返回: 注意力输出向量 (形状 1 x d_v), 注意力权重 (形状 1 x L_current)
    """
    # 1. 打分: 计算 Q_current 与 K_cache 的点积相似度
    # (1 x d_k) @ (d_k x L_current) -> (1 x L_current)
    scores = np.dot(Q_current, K_cache.T)

    # 2. 缩放: 防止数值过大导致 Softmax 梯度消失
    scaled_scores = scores / np.sqrt(d_k)

    # 3. 归一化: Softmax 得到注意力权重 (0到1之间，和为1)
    attention_weights = np.exp(scaled_scores - np.max(scaled_scores)) / np.sum(np.exp(scaled_scores - np.max(scaled_scores)), axis=-1, keepdims=True)

    # 4. 求和: 用注意力权重对 V_cache 进行加权求和
    # (1 x L_current) @ (L_current x d_v) -> (1 x d_v)
    attention_output = np.dot(attention_weights, V_cache)

    return attention_output, attention_weights

# --- 3. 推理生成模拟主流程 ---

def generate_text_with_kv_cache(user_prompt_text, max_new_tokens=5, mock_predicted_tokens=None):
    """
    模拟Transformer模型文本生成过程，带用户提示词输入和KV Cache。

    user_prompt_text: 用户的输入提示词字符串。
    max_new_tokens: 最多生成多少个新 Token。
    mock_predicted_tokens: (仅用于演示) 预设的模型预测结果列表，模拟模型生成。
                           如果为 None，则默认在生成max_new_tokens后生成<EOS>。
    """
    # 1. 将用户提示词转换为 Token IDs
    # 实际模型中，会有一个复杂的Tokenizer，这里简单地按空格分词并查找ID
    prompt_tokens = user_prompt_text.split()
    # 加上 <SOS> 作为序列的真正开始
    prompt_ids = [vocab["<SOS>"]] + [vocab.get(token, vocab["<PAD>"]) for token in prompt_tokens]

    generated_ids = list(prompt_ids) # 包含 <SOS> 和用户提示词
    output_text_tokens = [id_to_token[idx] for idx in generated_ids]

    # KV Cache 初始化为空
    kv_cache_K = np.empty((0, d_k))
    kv_cache_V = np.empty((0, d_v))

    print("=" * 60)
    print("--- 文本生成模拟开始 ---")
    print(f"用户提示词: '{user_prompt_text}'")
    print(f"初始模型输入 (Token IDs): {generated_ids} -> {' '.join(output_text_tokens)}")
    print("-" * 60)

    # --- 阶段 1: 处理初始提示词 (Prompt Processing / Pre-fill) ---
    # 这时模型会一次性计算整个提示词序列的K和V，并填充KV Cache。
    # 这一步是并行的，非常高效。
    print("\n[阶段 1] 处理用户提示词并填充 KV Cache...")
    # 获取整个提示词序列的 Embedding 矩阵 (L_prompt x d_model)
    prompt_embeddings = fixed_embedding_weights[prompt_ids]

    # 计算整个提示词序列的 K 和 V 矩阵
    # K_prompt = (L_prompt x d_model) @ (d_model x d_k) = (L_prompt x d_k)
    K_prompt = np.dot(prompt_embeddings, fixed_W_K)
    V_prompt = np.dot(prompt_embeddings, fixed_W_V)

    # 将提示词的 K 和 V 填充到 KV Cache
    kv_cache_K = np.vstack((kv_cache_K, K_prompt))
    kv_cache_V = np.vstack((kv_cache_V, V_prompt))

    print(f"  用户提示词处理完成。KV Cache 形状：K {kv_cache_K.shape}, V {kv_cache_V.shape}")
    print("-" * 60)

    # --- 阶段 2: 逐个 Token 生成 (Decoding) ---
    print("\n[阶段 2] 逐个 Token 生成模型回应...")
    for step in range(max_new_tokens):
        print(f"\n--- 第 {step + 1} 步：生成回应的下一个 Token ---")

        # 当前用于预测的 Token 总是整个序列的最后一个 Token。
        # 在这里，它要么是用户提示词的最后一个 Token，要么是模型上一步生成的新 Token。
        current_token_for_prediction_id = generated_ids[-1]
        current_token_for_prediction_text = id_to_token[current_token_for_prediction_id]
        print(f"  模型基于当前上下文的最后一个 Token '{current_token_for_prediction_text}' (ID: {current_token_for_prediction_id}) 进行预测。")

        # 3.1 获取当前预测 Token 的 Embedding
        # 形状 (1, d_model)
        current_token_embedding = fixed_embedding_weights[current_token_for_prediction_id].reshape(1, -1)
        print(f"  获取其 Embedding (形状: {current_token_embedding.shape})")

        # 3.2 **仅计算当前 Token 的 Query**
        # 形状 (1, d_k)
        Q_current = np.dot(current_token_embedding, fixed_W_Q)
        print(f"  计算其 Query Q (形状: {Q_current.shape})")

        # 3.3 计算注意力 (使用当前 Q 和 完整的 KV Cache)
        # 注意力的 K 和 V 来自于 KV Cache，包含了所有用户提示词和模型已生成的历史 Token 的信息
        attention_output, attention_weights = calculate_attention(Q_current, kv_cache_K, kv_cache_V)

        # 打印注意力权重，可以看到模型如何“关注”历史上下文
        print(f"  注意力权重 (关注整个历史序列，长度 {attention_weights.shape[1]}): {attention_weights[0].round(3)}")
        print(f"  注意力输出 (形状 {attention_output.shape}): {attention_output.sum(axis=1).round(3)}")

        # 3.4 模拟模型预测下一个 Token
        # 实际模型会根据 attention_output 和其他层（如前馈网络）的输出，
        # 预测下一个词的概率分布，然后进行采样或贪婪选择。
        if mock_predicted_tokens and step < len(mock_predicted_tokens):
            next_token_id = mock_predicted_tokens[step]
        else:
            # 如果没有预设预测，或者达到预设长度，就预测 <EOS>
            next_token_id = vocab["<EOS>"]

        next_token_text = id_to_token[next_token_id]
        print(f"  模型预测下一个 Token: '{next_token_text}'")

        # 3.5 停止条件检查
        if next_token_id == vocab["<EOS>"]:
            print(f"  遇到 <EOS> Token，生成结束。")
            break

        # 3.6 将新生成的 Token 添加到序列中 (作为新的上下文)
        generated_ids.append(next_token_id)
        output_text_tokens.append(next_token_text)

        # 3.7 **计算新生成 Token 的 K 和 V，并更新 KV Cache**
        # 针对刚刚预测出的 next_token_id，计算它的 K 和 V
        next_token_embedding = fixed_embedding_weights[next_token_id].reshape(1, -1)
        K_next_token = np.dot(next_token_embedding, fixed_W_K)
        V_next_token = np.dot(next_token_embedding, fixed_W_V)

        kv_cache_K = np.vstack((kv_cache_K, K_next_token))
        kv_cache_V = np.vstack((kv_cache_V, V_next_token))
        print(f"  将新 Token '{next_token_text}' 的 K/V 添加到缓存。")
        print(f"  KV Cache 形状更新为：K {kv_cache_K.shape}, V {kv_cache_V.shape}")

        print(f"  当前完整序列 (用于下一次预测): {' '.join(output_text_tokens)}")
        print("-" * 60)

    print("\n--- 文本生成完成 ---")
    # 过滤掉 <SOS> 标记，只显示用户输入和模型回应
    final_output_text = " ".join(output_text_tokens[1:]) # 跳过 <SOS>
    print(f"最终生成序列: {final_output_text}")
    print("=" * 60)

# --- 运行模拟 ---
# 模拟用户输入一个问题
user_input = "你 好 世界"
# 模拟模型希望生成的回应：“我 爱 北京”
mock_model_response_tokens = [vocab["我"], vocab["爱"], vocab["北京"], vocab["<EOS>"]]

generate_text_with_kv_cache(
    user_prompt_text=user_input,
    max_new_tokens=5, # 允许生成最多5个Token (包括可能的<EOS>)
    mock_predicted_tokens=mock_model_response_tokens
)

print("\n\n尝试另一个例子：")
generate_text_with_kv_cache(
    user_prompt_text="我 爱",
    max_new_tokens=2,
    mock_predicted_tokens=[vocab["你"], vocab["<EOS>"]] # 模拟生成“你”然后结束
)