File size: 2,591 Bytes
e791b16 c7770ff e791b16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import torch
from typing import List
def gpu_reward_simple(
generated_ids: torch.Tensor, # (batch, seq_len) on GPU
ground_truth_ids: torch.Tensor, # (batch, gt_len) on GPU
eos_token_id: int,
) -> torch.Tensor:
"""Simple GPU reward for token‑level subsequence matching.
Returns a tensor of shape (batch,) with 1.0 for a match and 0.0 otherwise.
"""
batch_size, seq_len = generated_ids.shape
rewards = torch.zeros(batch_size, device=generated_ids.device)
# locate EOS for each sequence
eos_mask = generated_ids == eos_token_id
# position of first EOS, or seq_len if none
first_eos = torch.where(
eos_mask.any(dim=1),
eos_mask.int().argmax(dim=1),
torch.full((batch_size,), seq_len, device=generated_ids.device, dtype=torch.long),
)
for i in range(batch_size):
gen = generated_ids[i, : first_eos[i]]
gt = ground_truth_ids[i]
gt_len = gt.shape[0]
if gt_len == 0 or gt_len > gen.shape[0]:
continue
# sliding window comparison
windows = gen.unfold(0, gt_len, 1) # (gen_len-gt_len+1, gt_len)
match = (windows == gt).all(dim=1).any()
rewards[i] = 1.0 if match else 0.0
return rewards
def tensor_native_reward(
generated_ids: torch.Tensor,
answer_patterns: List[torch.Tensor],
eos_token_id: int,
device: str = "cpu",
) -> torch.Tensor:
"""Fully batched tensor‑native reward.
- Trims each sequence at its first EOS token.
- For each batch element, checks whether the corresponding answer pattern
appears as a contiguous subsequence.
- Returns a (B,) float tensor of rewards (1.0 / 0.0).
"""
B, L = generated_ids.shape
device_obj = torch.device(device)
generated_ids = generated_ids.to(device_obj)
# Ensure answer patterns are on the same device
answer_patterns = [p.to(device_obj) for p in answer_patterns]
rewards = torch.zeros(B, device=device_obj)
# EOS handling
eos_mask = generated_ids == eos_token_id
first_eos = torch.where(
eos_mask.any(dim=1),
eos_mask.int().argmax(dim=1),
torch.full((B,), L, device=device_obj, dtype=torch.long),
)
for i in range(B):
seq = generated_ids[i, : first_eos[i]]
pattern = answer_patterns[i]
pat_len = pattern.shape[0]
if pat_len == 0 or pat_len > seq.shape[0]:
continue
windows = seq.unfold(0, pat_len, 1)
match = (windows == pattern).all(dim=1).any()
rewards[i] = 1.0 if match else 0.0
return rewards
|