Class FixedLogitsModel

tests/test_alignment_model_wrappers.py:20–40

class FixedLogitsModel(torch.nn.Module):
    def __init__(self, logits: torch.Tensor, *, output_kind: str = "tensor"):
        super().__init__()
        self.output_kind = output_kind
        self.weight = torch.nn.Parameter(torch.tensor(1.0))
        self.register_buffer("base_logits", logits.clone())

    def forward(self, input_ids: torch.Tensor, *, logit_bias: float = 0.0):
        del input_ids
        logits = self.base_logits * self.weight + float(logit_bias)
        if self.output_kind == "tensor":
            return logits
        if self.output_kind == "mapping":
            return {"logits": logits}
        if self.output_kind == "object":
            return ObjectOutput(logits)
        if self.output_kind == "tuple":
            return (logits, {"hidden_states": None})
        if self.output_kind == "loss_tuple":
            return (logits.sum(), logits, {"hidden_states": None})
        raise AssertionError(f"unknown output kind: {self.output_kind}")


@pytest.mark.parametrize("output_kind", ("tensor", "mapping", "object", "tuple"))