File size: 2,118 Bytes
dbc69f3 | 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 | from __future__ import annotations
import re
import torch
from transformers import PreTrainedTokenizerBase
# Inner reasoning body only (tags excluded from the masked span).
THINK_INNER_RE = re.compile(
r"<redacted_thinking>(.*?)</redacted_thinking>", re.DOTALL
)
def redacted_thinking_kl_scale(
completion_ids: torch.Tensor,
completion_mask: torch.Tensor,
tokenizer: PreTrainedTokenizerBase,
inner_kl_weight: float,
) -> torch.Tensor:
"""
Per completion token, return a multiplier for the KL term in GRPO.
- Default / outside the inner ``<redacted_thinking>...</redacted_thinking>`` body: 1.0
- Tokens whose character span overlaps the inner body: ``inner_kl_weight``
If ``inner_kl_weight == 1.0``, returns ones without decoding (fast path).
"""
device = completion_ids.device
bsz, seqlen = completion_ids.shape
scale = torch.ones((bsz, seqlen), device=device, dtype=torch.float32)
if inner_kl_weight == 1.0:
return scale
for b in range(bsz):
valid_len = int(completion_mask[b].sum().item())
if valid_len <= 0:
continue
row_ids = completion_ids[b, :valid_len].tolist()
text = tokenizer.decode(row_ids, skip_special_tokens=True)
match = THINK_INNER_RE.search(text)
if not match:
continue
inner_start, inner_end = match.span(1)
try:
enc = tokenizer(
text,
add_special_tokens=False,
return_offsets_mapping=True,
)
except TypeError:
continue
offsets = enc.offset_mapping
ids = enc.input_ids
if len(ids) != valid_len or len(offsets) != valid_len:
continue
for ti in range(min(valid_len, seqlen)):
if int(completion_mask[b, ti].item()) == 0:
continue
cs, ce = offsets[ti]
if ce <= cs:
continue
overlaps = cs < inner_end and ce > inner_start
if overlaps:
scale[b, ti] = float(inner_kl_weight)
return scale
|