noao-chat-6-训练评估指南

@torch.no_grad()
def evaluate_bpb(model, batches, steps, token_bytes):
    """
    计算 Bits Per Byte
    
    Args:
        model: 要评估的模型
        batches: 数据迭代器
        steps: 评估步数
        token_bytes: (vocab_size,) 张量，每个 token 的字节数
    
    Returns:
        float: BPB 值
    """
    total_nats = 0.0  # 总自然对数损失
    total_bytes = 0   # 总字节数
    
    for _ in range(steps):
        x, y = next(batches)
        
        # 计算每个位置的损失（不降维）
        loss2d = model(x, y, loss_reduction='none')  # (B, T)
        
        # 获取每个 token 的字节数
        valid = y >= 0  # 排除 ignore_index=-1
        y_safe = torch.where(valid, y, torch.zeros_like(y))
        num_bytes = token_bytes[y_safe]  # (B, T)
        
        # 累积
        total_nats += (loss2d * (num_bytes > 0)).sum()
        total_bytes += num_bytes.sum()
    
    # 转换：nats → bits, 归一化
    bpb = total_nats / (total_bytes * math.log(2))
    return bpb

CORE_TASKS = [
    # 1. 多选题 (Multiple Choice)
    "arc_easy", "arc_challenge",
    "hellaswag", "piqa", "winogrande",
    
    # 2. 模式匹配 (Schema)
    "copa", "openbookqa",
    
    # 3. 语言建模 (Language Modeling)
    "lambada", "wikitext"
]

core_metric = mean([accuracy(task) for task in CORE_TASKS])

def evaluate_task(model, tokenizer, data, device, task_meta):
    correct = torch.zeros(len(data))
    
    for idx in range(len(data)):
        item = data[idx]
        
        # 1. 采样 few-shot 示例
        fewshot_examples = sample_fewshot(data, idx, num_fewshot)
        
        # 2. 渲染 prompt
        prompts = render_prompts(item, fewshot_examples, task_type)
        
        # 3. Tokenize
        tokens = tokenizer(prompts)
        
        # 4. 前向传播，计算损失
        losses = model(tokens)
        
        # 5. 判断正确性
        if task_type == 'multiple_choice':
            # 选择损失最小的选项
            pred_idx = losses.argmin()
            is_correct = (pred_idx == item['gold'])
        elif task_type == 'language_modeling':
            # 检查生成是否匹配
            is_correct = (predictions == targets).all()
        
        correct[idx] = is_correct
    
    return correct.mean()

accuracy = correct_count / total_count

def evaluate_pass_at_k(model, task, k=8):
    """
    评估 Pass@k
    
    Args:
        model: 模型
        task: 任务对象
        k: 采样次数
    
    Returns:
        float: Pass@k 准确率
    """
    correct = 0
    total = 0
    
    for question in task:
        # 生成 k 个答案
        answers = model.generate(
            question,
            num_samples=k,
            temperature=1.0  # 保持随机性
        )
        
        # 评估每个答案
        results = [task.evaluate(question, ans) for ans in answers]
        
        # 至少一个正确 → 算作 Pass
        passed = any(results)
        
        correct += int(passed)
        total += 1
    
    return correct / total

# SFT 模型
Pass@1 = 0.45   # 第一次尝试的准确率
Pass@8 = 0.72   # 给 8 次机会的准确率

# RL 模型（训练后）
Pass@1 = 0.53   # 提升 8%
Pass@8 = 0.78   # 提升 6%

# 问题
question = "What is the primary function of mitochondria?"

# 选项
choices = [
    "Protein synthesis",
    "Energy production",      # ← 正确答案
    "DNA replication",
    "Cell division"
]

# 用户消息（渲染）
user_message = """
What is the primary function of mitochondria?
A) Protein synthesis
B) Energy production
C) DNA replication
D) Cell division
"""

# 助手回答
assistant_message = "B"

def evaluate_mmlu(model, tokenizer):
    for problem in mmlu_dataset:
        # 1. 构造 4 个完整 prompt（每个选项一个）
        prompts = [
            f"{question}\n{choices}\nAnswer: A",
            f"{question}\n{choices}\nAnswer: B",
            f"{question}\n{choices}\nAnswer: C",
            f"{question}\n{choices}\nAnswer: D",
        ]
        
        # 2. 计算每个选项的损失
        losses = [model.compute_loss(prompt) for prompt in prompts]
        
        # 3. 选择损失最小的
        pred_letter = letters[losses.argmin()]
        
        # 4. 比较
        is_correct = (pred_letter == correct_letter)

# 问题
question = """
Weng earns $12 an hour for babysitting. 
Yesterday, she just did 50 minutes of babysitting. 
How much did she earn?
"""

# 标准答案（带计算过程）
answer = """
Weng earns 12/60 = $<<12/60=0.2>>0.2 per minute.
Working 50 minutes, she earned 0.2 x 50 = $<<0.2*50=10>>10.
#### 10
"""

def extract_answer(completion):
    """提取 #### 后的数字"""
    match = re.search(r"#### (\-?[0-9\.\,]+)", completion)
    if match:
        return match.group(1).replace(",", "")
    return None

def evaluate(conversation, completion):
    # 提取参考答案
    ref_answer = extract_answer(conversation['messages'][-1])
    
    # 提取生成的答案
    pred_answer = extract_answer(completion)
    
    # 比较
    return float(ref_answer == pred_answer)

# Prompt（函数签名）
prompt = """
def has_close_elements(numbers: List[float], threshold: float) -> bool:
    \"\"\" Check if in given list of numbers, are any two numbers closer
    to each other than given threshold.
    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)
    False
    >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)
    True
    \"\"\"
"""

# 标准答案（函数实现）
canonical_solution = """
    for idx, elem in enumerate(numbers):
        for idx2, elem2 in enumerate(numbers):
            if idx != idx2:
                distance = abs(elem - elem2)
                if distance < threshold:
                    return True
    return False
"""

def evaluate(conversation, completion):
    # 1. 提取代码
    code = extract_program(completion)
    
    # 2. 构造完整程序
    program = f"""
{imports}

{code}

{test_cases}

check({entry_point})
"""
    
    # 3. 执行代码
    result = execute_code(program, timeout=5)
    
    # 4. 返回是否成功
    return result.success

# 用户问题
question = "How many r are in strawberry?"

# 助手回答（带推理）
answer = """
Let me spell out the word:
s-t-r-a-w-b-e-r-r-y

Now let me count the letter 'r':
<|python_start|>"strawberry".count("r")<|python_end|>
<|output_start|>3<|output_end|>

#### 3
"""

def evaluate(conversation, completion):
    # 提取参考答案
    ref = extract_answer(conversation['messages'][-1])
    
    # 提取生成的答案
    pred = extract_answer(completion)
    
    # 比较数字
    return float(ref == pred)

# scripts/base_loss.py
device_batch_size = 64
sequence_len = 2048
split_tokens = 20 * 524288  # 评估 20M tokens

for split_name in ['train', 'val']:
    # 创建数据加载器
    loader = tokenizing_distributed_data_loader(
        device_batch_size, 
        sequence_len, 
        split_name
    )
    
    # 计算步数
    steps = split_tokens // (device_batch_size * sequence_len * world_size)
    
    # 评估
    bpb = evaluate_bpb(model, loader, steps, token_bytes)
    
    print(f"{split_name} BPB: {bpb:.4f}")

# scripts/base_eval.py
def evaluate_model(model, tokenizer, device, max_per_task=-1):
    # 1. 加载 CORE 配置
    tasks = load_core_config()
    
    results = {}
    for task in tasks:
        # 2. 加载任务数据
        data = load_task_data(task)
        
        # 3. 评估
        accuracy = evaluate_task(
            model, 
            tokenizer, 
            data, 
            device, 
            task_meta
        )
        
        results[task['label']] = accuracy
    
    # 4. 计算 CORE Metric
    core_metric = mean(results.values())
    
    return core_metric, results

# 在训练脚本中
if step % sample_every == 0:
    model.eval()
    
    # 生成样本
    prompt = "The capital of France is"
    tokens = tokenizer.encode(prompt)
    
    for token in model.generate(tokens, max_tokens=50):
        print(tokenizer.decode([token]), end='')
    
    print()
    model.train()

def run_categorical_eval(task_object, tokenizer, model, batch_size):
    """
    分类任务评估（不需要采样）
    """
    num_correct = 0
    total = 0
    
    # 按批次处理
    for batch_start in range(0, len(task_object), batch_size):
        batch = task_object[batch_start:batch_start + batch_size]
        
        # 1. 准备所有选项的 prompt
        all_prompts = []
        for problem in batch:
            # 每个问题有多个选项
            for letter in problem['letters']:
                prompt = tokenizer.render_with_answer(
                    problem, 
                    letter
                )
                all_prompts.append(prompt)
        
        # 2. Tokenize 并填充
        tokens = tokenizer(all_prompts)
        input_ids = pad_sequences(tokens)
        
        # 3. 前向传播，计算每个选项的 logits
        logits = model(input_ids)
        
        # 4. 计算每个选项在答案位置的 log prob
        # （细节：找到答案 token 的位置）
        answer_logprobs = extract_answer_logprobs(logits, tokens)
        
        # 5. 选择 logprob 最高的
        num_choices = len(problem['letters'])
        for i, problem in enumerate(batch):
            choice_logprobs = answer_logprobs[
                i * num_choices : (i+1) * num_choices
            ]
            pred_idx = choice_logprobs.argmax()
            correct_idx = problem['letters'].index(
                problem['correct_answer']
            )
            
            if pred_idx == correct_idx:
                num_correct += 1
            total += 1
    
    return num_correct / total

def run_generative_eval(task_object, tokenizer, model, engine, 
                        num_samples, max_new_tokens, temperature, top_k):
    """
    生成任务评估（需要采样）
    """
    num_passed = 0
    total = 0
    
    for problem in task_object:
        # 1. Tokenize prompt
        encoded_prompt = tokenizer.render_for_completion(problem)
        
        # 2. 生成多个答案
        results, _ = engine.generate_batch(
            encoded_prompt,
            num_samples=num_samples,    # 例如 8
            max_tokens=max_new_tokens,  # 例如 512
            temperature=temperature,    # 例如 0.8
            top_k=top_k,                # 例如 50
        )
        
        # 3. 解码
        prefix_length = len(encoded_prompt)
        completions = [
            tokenizer.decode(result[prefix_length:])
            for result in results
        ]
        
        # 4. 评估每个答案
        outcomes = [
            task_object.evaluate(problem, completion)
            for completion in completions
        ]
        
        # 5. Pass@k：至少一个正确
        passed = any(outcomes)
        
        num_passed += int(passed)
        total += 1
    
    return num_passed / total

# scripts/chat_rl.py
for step in range(num_steps):
    # 训练...
    
    # 定期评估
    if step % eval_every == 0:
        model.eval()
        
        # 评估多个 k 值
        pass_at_k = {}
        for k in [1, 4, 8, 16]:
            accuracy = evaluate_pass_at_k(
                model, 
                val_task, 
                k=k,
                max_problems=eval_examples
            )
            pass_at_k[f"pass@{k}"] = accuracy
        
        print(f"Step {step} | Pass@1: {pass_at_k['pass@1']:.3f} | Pass@8: {pass_at_k['pass@8']:.3f}")
        
        model.train()

def evaluate_example(idx, model, tokenizer, data, task_meta):
    item = data[idx]
    num_fewshot = task_meta['num_fewshot']  # 例如 5
    
    # 从数据集中随机采样 few-shot 示例（排除当前样本）
    fewshot_examples = []
    if num_fewshot > 0:
        rng = random.Random(1234 + idx)  # 固定种子保证可复现
        available_indices = [i for i in range(len(data)) if i != idx]
        fewshot_indices = rng.sample(available_indices, num_fewshot)
        fewshot_examples = [data[i] for i in fewshot_indices]
    
    # 渲染 prompt
    prompt = render_prompt_with_fewshot(item, fewshot_examples)
    
    # ...评估

def stack_sequences(tokens, pad_token_id):
    """
    将多个序列堆叠成批次
    
    Args:
        tokens: List[List[int]]，每个序列的 token IDs
        pad_token_id: 填充 token
    
    Returns:
        input_ids: (B, max_len) 张量
    """
    bsz = len(tokens)
    seq_len = max(len(x) for x in tokens)
    
    # 创建填充后的张量
    input_ids = torch.full(
        (bsz, seq_len), 
        pad_token_id, 
        dtype=torch.long
    )
    
    # 填入实际 tokens
    for i, x in enumerate(tokens):
        input_ids[i, :len(x)] = torch.tensor(x, dtype=torch.long)
    
    return input_ids

tokens = [
    [1, 2, 3, 4],         # 长度 4
    [1, 2, 3, 4, 5, 6],   # 长度 6
    [1, 2],               # 长度 2
]

# Padding
input_ids = stack_sequences(tokens, pad_token_id=0)
# tensor([
#     [1, 2, 3, 4, 0, 0],
#     [1, 2, 3, 4, 5, 6],
#     [1, 2, 0, 0, 0, 0],
# ])

def find_common_length(token_sequences, direction='left'):
    """
    找到所有序列的公共前缀长度
    """
    min_len = min(len(seq) for seq in token_sequences)
    
    for i in range(min_len):
        token = token_sequences[0][i]
        if not all(seq[i] == token for seq in token_sequences):
            return i  # 第一个不同的位置
    
    return min_len

# 示例
prompts = [
    "What is the capital of France? A) London",
    "What is the capital of France? B) Paris",
    "What is the capital of France? C) Berlin",
    "What is the capital of France? D) Madrid",
]

tokens = tokenizer(prompts)
answer_start = find_common_length(tokens)
# "What is the capital of France? " 的长度
# 答案从这里开始

# 方向='right'
answer_length = find_common_length(tokens, direction='right')
# 从末尾往前数 answer_length 个 token

tokens_without = tokenizer(prompt_without)  # 不含答案
tokens_with = tokenizer(prompt_with)        # 含答案

# 答案位置
answer_start = len(tokens_without)
answer_end = len(tokens_with)

def calculate_answer_loss(logits, targets, start_idx, end_idx):
    """
    只计算 [start_idx, end_idx) 范围内的损失
    """
    # logits: (B, T, vocab_size)
    # targets: (B, T)
    
    # 创建 mask
    mask = torch.zeros_like(targets, dtype=torch.bool)
    mask[:, start_idx:end_idx] = True
    
    # 计算所有位置的损失
    losses = F.cross_entropy(
        logits.view(-1, vocab_size),
        targets.view(-1),
        reduction='none'
    ).view(B, T)
    
    # 只保留答案部分
    answer_losses = losses * mask
    
    # 平均
    mean_loss = answer_losses.sum() / mask.sum()
    
    return mean_loss

# scripts/chat_rl.py
eval_every = 60
eval_examples = 400

for step in range(num_steps):
    # ... 训练 ...
    
    if step % eval_every == 0:
        # 评估多个 k 值
        for k in [1, 4, 8, 16]:
            accuracy = evaluate_pass_at_k(
                model, 
                val_task, 
                k=k,
                max_problems=eval_examples
            )
            wandb.log({f"pass@{k}": accuracy})

# 开发阶段：快速测试
max_per_task = 100  # 只评估 100 个样本

# 生产阶段：完整评估
max_per_task = -1   # 评估所有样本

# 分类任务：可以批量处理
batch_size = 32  # 同时处理 32 个问题

# 生成任务：逐个处理（因为长度不同）
batch_size = 1

# 缓存 few-shot 示例
@lru_cache(maxsize=1000)
def get_fewshot_examples(task_name, example_idx, num_fewshot):
    # ...
    return fewshot_examples

# 调试时打印实际的 prompt
if idx == 0:  # 只打印第一个样本
    print("=" * 50)
    print("Prompt:")
    print(prompt)
    print("=" * 50)

# 验证答案位置是否正确
prompt_tokens = tokenizer(prompt)
print(f"Total length: {len(prompt_tokens)}")
print(f"Answer start: {answer_start_idx}")
print(f"Answer tokens: {prompt_tokens[answer_start_idx:]}")

# 1. 先在单个样本上测试
result = evaluate_example(0, model, tokenizer, data, task_meta)
print(f"First example: {result}")

# 2. 再在小批次上测试
results = [evaluate_example(i, ...) for i in range(10)]
print(f"First 10 examples: {sum(results) / len(results)}")

# 3. 最后完整评估
accuracy = evaluate_task(model, tokenizer, data, device, task_meta)

# 打印并检查
print(f"Prompt: {tokenizer.decode(tokens[:answer_start])}")
print(f"Answer: {tokenizer.decode(tokens[answer_start:])}")

# 使用 attention_mask
loss = F.cross_entropy(
    logits.view(-1, vocab_size),
    targets.view(-1),
    ignore_index=pad_token_id,  # ← 关键
    reduction='none'
)

# Pass@k 必须用随机采样
temperature = 1.0  # 或 0.7, 0.8
top_k = 50