Class PolicyModelWrapper

rl_engine/alignment/model_wrappers.py:77–110

Standard adapter for the trainable policy model used by RL losses.

class PolicyModelWrapper(torch.nn.Module):
    """Standard adapter for the trainable policy model used by RL losses."""

    def __init__(self, model: torch.nn.Module):
        super().__init__()
        self.model = model

    def forward(self, input_ids: torch.Tensor, **model_kwargs: Any) -> Any:
        return self.model(input_ids, **model_kwargs)

    def forward_logits(self, input_ids: torch.Tensor, **model_kwargs: Any) -> torch.Tensor:
        return extract_logits(self.forward(input_ids, **model_kwargs))

    def selected_logprobs(
        self,
        input_ids: torch.Tensor,
        token_ids: torch.Tensor,
        *,
        mask: Optional[torch.Tensor] = None,
        logits_start: Optional[int] = None,
        logits_end: Optional[int] = None,
        temperature: float = 1.0,
        output_dtype: torch.dtype = torch.float32,
        **model_kwargs: Any,
    ) -> torch.Tensor:
        logits = self.forward_logits(input_ids, **model_kwargs)
        logits = _slice_logits(logits, logits_start=logits_start, logits_end=logits_end)
        return selected_logprobs_reference(
            logits,
            token_ids,
            mask=mask,
            temperature=temperature,
            output_dtype=output_dtype,
        )


class ReferenceModelWrapper(PolicyModelWrapper):