SupplementaryMaterial_Code/standard/VLMEvalKit/test2.py

import torch

# 定义你的函数（直接复制）
def filter_logits_with_candidates(logits, candidate_indices):
    """
    高效版本：使用向量化索引替代for循环，提升计算效率
    """
    # 初始化全为-inf的张量
    filtered_logits = torch.full_like(logits, float('-inf'))
    
    # 获取batch_size和候选数量
    batch_size = logits.size(0)
    num_candidates = candidate_indices.size(1)
    
    # 构造batch维度的索引
    batch_indices = torch.arange(batch_size, device=logits.device)[:, None].expand(-1, num_candidates)
    print(batch_indices)
    # 高级索引一次性更新
    filtered_logits[batch_indices, candidate_indices] = logits[batch_indices, candidate_indices]
    
    return filtered_logits

# ---------------------- 输入示例构造 ----------------------
# 1. 构造logits：shape=(batch_size, vocab_size)
# batch_size=2（2个样本），vocab_size=10（模拟10个候选词的logits）
logits = torch.tensor([
    [1.2, 3.4, 0.5, 2.1, 4.3, 1.8, 0.9, 3.7, 2.5, 1.1],  # 第1个样本的logits
    [2.8, 0.7, 3.1, 1.5, 4.0, 2.2, 3.5, 0.3, 1.9, 2.6]   # 第2个样本的logits
])

# 2. 构造候选索引：shape=(batch_size, num_candidates)
# 每个样本选3个候选词的索引（索引范围必须在[0, vocab_size-1]，即0-9）
candidate_indices = torch.tensor([
    [2, 5, 7],  # 第1个样本的候选索引：对应logits[0,2]=0.5、logits[0,5]=1.8、logits[0,7]=3.7
    [1, 4, 8]   # 第2个样本的候选索引：对应logits[1,1]=0.7、logits[1,4]=4.0、logits[1,8]=1.9
])

# ---------------------- 函数调用与验证 ----------------------
filtered_logits = filter_logits_with_candidates(logits, candidate_indices)

# 打印输入和输出，对比验证
print("原始logits（shape={}）：".format(logits.shape))
print(logits)
print("\n候选索引（shape={}）：".format(candidate_indices.shape))
print(candidate_indices)
print("\n过滤后logits（仅候选索引位置保留原值，其余为-inf）：")
print(filtered_logits)