# -*- coding: utf-8 -*-
"""
僅此文件允許考生修改：
- 請在下列函數的函數體內完成實現。
- 不要改動函數名與參數簽名。
- 你可以新增少量輔助函數。
"""

from typing import List
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM


# ============================================================
# 第一部分：Prompt 定義
# ============================================================
def build_system_prompt() -> str:
    """
    考生實現：定義 system prompt
    - 返回一個 system prompt，要求模型以"[Answer]: xxxx"的格式給出最終數值。
    
    System Prompt 的作用：
    1. 定義模型的身份和角色（數學助手）
    2. 指導模型的推理過程（逐步分析、解題思路）
    3. 規范模型的輸出格式（[Answer]: 數值）
    
    為什么要用 [Answer]: 格式？
    - 便于評測程序從模型輸出中提取正確答案
    - 統一答案格式，避免模型輸出過多解釋性文字
    - 符合測試數據的標準格式要求
    """
    # ======== 考生實現區域（可修改） ========
    
    # System Prompt 定義：指導模型如何回答數學問題
    # 關鍵點：必須包含 [Answer]: 格式的要求，這是評測程序提取答案的標準
    system_prompt = """你是一個數學助手，擅長解決各種數學問題。請仔細分析問題，逐步推理，并最終以 [Answer]: 數值 的格式給出答案。

要求：
1. 仔細閱讀題目，理解題意
2. 逐步分析解題思路
3. 進行必要的計算
4. 最終答案必須以 [Answer]: 數值 的格式給出

請開始解題。"""
    
    return system_prompt
    
    # ======== 考生實現區域（可修改） ========


# ============================================================
# 第二部分：模板拼裝
# ============================================================
def apply_chat_template_single(
    tokenizer: AutoTokenizer,
    system_prompt: str,
    problem: str,
) -> str:
    """
    考生實現：將單個問題轉換為模型輸入文本
    - 使用 tokenizer.apply_chat_template 構造對話
    - 返回拼裝好的文本字符串
    
    參數說明：
        tokenizer: Qwen3 模型的 tokenizer，包含特殊 token 和模板定義
        system_prompt: 系統提示詞，定義模型角色和輸出格式
        problem: 用戶問題（如 "12 + 35"）
    
    返回值：
        完整格式化后的對話文本，包含 system、user、assistant 標記
    
    Chat Template 的作用：
        transformers 庫的 tokenizer.apply_chat_template 會根據模型類型自動構造
        符合該模型預期的對話格式。對于 Qwen 模型，它會把 messages 列表轉換為：
        <|system|>你是一個數學助手...<|user|>12 + 35<|assistant|>
        這樣的格式。
    """
    # ======== 考生實現區域（可修改） ========
    
    # 構造消息列表：包含系統提示和用戶問題
    # 消息格式：每個消息是一個字典，包含 role（角色）和 content（內容）
    # - "system": 定義模型行為
    # - "user": 用戶輸入的問題
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": problem},
    ]
    
    # 應用聊天模板，將消息列表轉換為模型可理解的文本格式
    rendered = tokenizer.apply_chat_template(
        messages,                      # 消息列表
        tokenize=False,                # 不進行 tokenize，返回原始字符串而不是 token IDs
                                      # 這樣可以在調試時看到完整的文本內容
        add_generation_prompt=True,    # 添加 <|assistant|> 標記，引導模型生成回復
                                      # 如果沒有這個標記，模型不知道要在哪里開始生成
        enable_thinking=True,          # 啟用思考標簽支持
                                      # 允許模型輸出 <think>...</think>
                                      # 模型會在 <> 標簽內進行推理，評測程序會自動移除這部分內容
    )
    return rendered
    
    # ======== 考生實現區域（可修改） ========


# ============================================================
# 第三部分：核心推理實現
# ============================================================
def generate_single(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    rendered_text: str,
    max_new_tokens: int,
    do_sample: bool,
) -> torch.Tensor:
    """
    考生實現：單條推理
    - 將文本 tokenize 后送入模型生成
    - 返回包含輸入和輸出的完整 token 序列
    
    參數說明：
        model: Qwen3 預訓練模型，用于文本生成
        tokenizer: 用于將文本轉換為 token IDs
        rendered_text: 已格式化的完整對話文本（包含 system、user、assistant 標記）
        max_new_tokens: 生成的最大新 token 數量（不包括輸入長度）
        do_sample: 是否使用采樣（True=隨機采樣，False=貪心解碼）
    
    返回值：
        torch.Tensor，形狀為 [1, seq_len]，包含輸入+新生成的完整序列
    
    注意事項：
        1. 返回的是包含輸入的完整序列（evaluate.py 會提取新生成的部分）
        2. 使用 .to(model.device) 確保輸入在正確的設備上（CPU/GPU/NPU）
        3. 使用 model.generate() 的默認參數，模型會自動處理特殊 token
    """
    # ======== 考生實現區域（可修改） ========
    
    # 步驟 1：將文本編碼為 token IDs
    # - tokenizer.encode() 將字符串轉換為整數列表
    # - return_tensors="pt" 返回 PyTorch 張量
    # - .to(model.device) 將張量移動到模型所在的設備（CPU/GPU/NPU）
    inputs = tokenizer.encode(rendered_text, return_tensors="pt").to(model.device)
    
    # 步驟 2：使用模型生成文本
    # - model.generate() 會自動生成新的 token
    # - 返回的 outputs 包含完整的序列：輸入 + 新生成的 token
    # - 評測程序會從 outputs 中提取新生成的部分
    outputs = model.generate(
        inputs,                        # 輸入 token IDs
        max_new_tokens=max_new_tokens, # 最多生成多少個新 token
        do_sample=do_sample,          # 采樣模式：True=隨機采樣（更靈活），False=貪心（更確定）
    )
    return outputs
    # ======== 考生實現區域（可修改） ========


def generate_batch(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    rendered_texts: List[str],
    max_new_tokens: int,
    do_sample: bool,
) -> List[torch.Tensor]:
    """
    考生實現：批量推理
    - 一次處理多個問題，提高效率
    - 返回所有批次的輸出列表
    
    參數說明：
        model: Qwen3 預訓練模型
        tokenizer: 用于編碼文本
        rendered_texts: 多個已格式化的對話文本列表
        max_new_tokens: 生成的最大新 token 數量
        do_sample: 是否使用采樣
    
    返回值：
        List[torch.Tensor]，每個元素是一個批次的輸出
        批次形狀為 [batch_size, seq_len]
    
    批量處理的優勢：
        1. 并行處理：GPU 可以同時處理多條數據
        2. 減少數據傳輸：一次 GPU 調用處理多條
        3. 提高吞吐量：相比逐條處理能獲得 3+ 倍加速
    
    實現要點：
        1. 使用循環分批處理，避免一次性處理過多數據導致內存溢出
        2. 使用 padding=True 確保同一批次內序列長度一致
        3. 使用 attention_mask 告訴模型哪些 token 是 padding，哪些是真實內容
        4. 使用 torch.no_grad() 節省顯存（推理時不需要梯度）
    """
    # ======== 考生實現區域（可修改） ========
    
    all_outputs = []  # 存儲所有批次的輸出
    batch_size: int = 4  # 批次大小：一次處理 4 條數據
                        # 可以調整：太小效率低，太大可能內存不足

    # 循環處理：每次處理一個批次
    # range(0, len(rendered_texts), batch_size) 生成：0, 4, 8, 12, ...
    for i in range(0, len(rendered_texts), batch_size):
        # 提取當前批次：rendered_texts[i:i+batch_size]
        # 例如：如果 i=0, batch_size=4，則提取第 0, 1, 2, 3 條
        batch_texts = rendered_texts[i:i + batch_size]
        
        # 步驟 1：對批次進行編碼和填充
        # tokenizer() 可以將多條文本同時編碼，并自動填充到相同長度
        # - batch_texts: 多條文本（列表）
        # - padding=True: 短序列會填充到和最長序列一樣長
        # - truncation=True: 超長序列會被截斷
        # - 返回的 inputs 包含 input_ids 和 attention_mask
        inputs = tokenizer(
            batch_texts,               # 多條文本（如 ["問題1", "問題2", "問題3", "問題4"]）
            return_tensors="pt",        # 返回 PyTorch 張量
            padding=True,              # 自動填充到相同長度
            truncation=True            # 超長序列截斷（防止溢出）
        ).to(model.device)             # 移動到模型設備
        
        # 步驟 2：批量生成文本
        # 使用 torch.no_grad() 禁用梯度計算，節省顯存并提高速度
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids,        # 編碼后的 token IDs，形狀 [batch_size, seq_len]
                attention_mask=inputs.attention_mask,  # 注意力掩碼，標記哪些是真實內容（1），哪些是填充（0）
                max_new_tokens=max_new_tokens,  # 最多生成多少個新 token
                do_sample=do_sample,     # 采樣模式
                pad_token_id=tokenizer.eos_token_id  # 指定填充 token 的 ID
                # 注意：這里沒有使用 pad_token_id=tokenizer.pad_token_id
                # 因為 Qwen 模型可能沒有 pad_token，使用 eos_token_id 作為替代
            )
        
        # 將當前批次的輸出添加到總結果中
        all_outputs.append(outputs)

    # 返回所有批次的輸出列表
    # 例如：如果 10 條數據分成 3 批，則返回 [batch1_output, batch2_output, batch3_output]
    # evaluate.py 會將這些批次的結果合并到一起
    return all_outputs
    
    # ======== 考生實現區域（可修改） ========

# -*- coding: utf-8 -*-
"""
僅此文件允許考生修改：
- 請在下列函數的函數體內完成實現。
- 不要改動函數名與參數簽名。
- 你可以新增少量輔助函數。
"""

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel


# ============================================================
# 相似度計算函數
# ============================================================
def student_compute_similarity(text1: str, text2: str, model: AutoModel, tokenizer: AutoTokenizer) -> float:
    """
    考生實現：計算兩個文本之間的相似度
    
    參數：
        text1: 第一個文本字符串（如 "12 + 35"）
        text2: 第二個文本字符串（如 "35 + 12"）
        model: 預加載的 Qwen3-Embedding 模型（評測程序提供）
        tokenizer: 預加載的 tokenizer（評測程序提供）
    
    返回：
        相似度值（0.0 到 1.0 之間的浮點數）
        1.0 表示完全相同，0.0 表示完全不同
    
    要求：
        - 使用傳入的 model 和 tokenizer，不要自己加載模型
        - 實現 last-token pooling（取最后一個 token 的隱藏狀態作為句子表示）
        - 必須使用左側 padding（padding_side="left"）
        - L2 歸一化（將向量標準化為單位向量）
        - 計算余弦相似度（歸一化向量的點積）
        - 不得使用 sentence_transformers
    
    相似度計算原理：
        1. 將文本編碼為向量（embedding）
        2. Last-token pooling：取最后一個 token 的隱藏狀態作為句子表示
        3. L2 歸一化：將向量標準化為單位向量
        4. 余弦相似度：計算兩個歸一化向量的點積（cos(θ) 的值域為 [-1, 1]）
        
        為什么使用 last-token pooling？
        - Embedding 模型在訓練時，最后一個 token 包含了整個句子的語義信息
        - 相比平均 pooling 或 CLS token，last-token 能更好地捕捉句子級語義
    """
    # ======== 考生實現區域（可修改） ========
    
    # 步驟 1：獲取模型所在的設備（CPU/GPU/NPU）
    # next(model.parameters()) 獲取模型的第一個參數，.device 獲取其設備
    # 這樣可以確保輸入張量與模型在同一設備上
    device = next(model.parameters()).device
    
    # 步驟 2：準備輸入文本列表
    # 將兩個文本放入列表，以便批量處理
    texts = [text1, text2]
    
    # 步驟 3：使用 tokenizer 編碼文本
    # padding=True: 將短序列填充到最長序列的長度
    # truncation=True: 超長序列截斷到 max_length
    # max_length=512: 最大序列長度（典型的 BERT 等模型長度）
    # return_tensors="pt": 返回 PyTorch 張量
    # padding_side="left": 在左側填充（重要！）
    #   為什么使用左側 padding？
    #   - 因為使用 last-token pooling，我們需要保證最后一個 token 是真實的文本 token
    #   - 如果在右側 padding，最后一個 token 會是填充 token，失去語義信息
    inputs = tokenizer(
        texts,                      # 文本列表：["12 + 35", "35 + 12"]
        padding=True,              # 自動填充到相同長度
        truncation=True,            # 超長序列截斷
        max_length=512,            # 最大長度
        return_tensors="pt",        # 返回 PyTorch 張量
        padding_side="left"         # 關鍵：左側填充
    ).to(device)                    # 移動到模型設備
    
    # 步驟 4：獲取模型輸出的嵌入向量
    # 使用 torch.no_grad() 禁用梯度計算，節省顯存并提高速度
    with torch.no_grad():
        # 模型前向傳播：輸入 token IDs，輸出隱藏狀態
        # outputs.last_hidden_state 形狀: [batch_size=2, seq_len, hidden_dim]
        outputs = model(**inputs)
        
        # Last-token pooling：取最后一個 token 的隱藏狀態作為句子表示
        # last_hidden_state 形狀: [2, seq_len, hidden_dim]
        # 例如：[[[token1], [token2], [token3]], [[token4], [token5], [token6]]]
        # [:, -1, :] 表示取每一行的最后一個 token
        # embeddings 形狀: [2, hidden_dim]
        # 例如：[[token3], [token6]]
        last_hidden_state = outputs.last_hidden_state
        embeddings = last_hidden_state[:, -1, :]  # 索引 -1 表示最后一個元素
        
        # 步驟 5：L2 歸一化
        # 將向量標準化為單位向量（長度為 1）
        # 公式：v_normalized = v / ||v||
        # F.normalize(..., p=2, dim=1):
        #   - p=2: 使用 L2 范數（歐氏距離）
        #   - dim=1: 在 hidden_dim 維度上歸一化
        # 歸一化后的向量滿足：||v|| = 1
        # embeddings 形狀保持不變: [2, hidden_dim]
        embeddings = F.normalize(embeddings, p=2, dim=1)
        
        # 步驟 6：計算余弦相似度
        # 對于歸一化的向量 a, b，余弦相似度 = a · b（點積）
        # 如果 a 和 b 都是單位向量，則 cos(θ) = a · b
        # - 如果 a 和 b 完全相同：a · b = 1
        # - 如果 a 和 b 垂直：a · b = 0
        # - 如果 a 和 b 相反：a · b = -1
        # torch.dot(a, b) 計算向量點積
        # .item() 將標量張量轉換為 Python float
        # embeddings[0] 是 text1 的嵌入，embeddings[1] 是 text2 的嵌入
        similarity = torch.dot(embeddings[0], embeddings[1]).item()
        # similarity 的范圍：[0, 1]，因為所有向量都是正方向（經過歸一化后的嵌入向量）
    
    return similarity
    
    # ======== 考生實現區域（可修改） ========


def compute_similarity(text1: str, text2: str, model: AutoModel, tokenizer: AutoTokenizer) -> float:
    """
    評測程序調用的接口函數
    """
    return student_compute_similarity(text1, text2, model, tokenizer)