| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | |
| 8 | @dataclass(frozen=True) |
| 9 | class Oracle: |
| 10 | token_ids: tuple[int, ...] |
| 11 | text: str |
| 12 | sampled_token_logprob: float |
| 13 | |
| 14 | def __post_init__(self) -> None: |
| 15 | if len(self.token_ids) != 1: |
| 16 | raise ValueError("snapshot oracle must contain exactly one token") |
| 17 | value = self.sampled_token_logprob |
| 18 | if ( |
| 19 | isinstance(value, bool) |
| 20 | or not isinstance(value, (int, float)) |
| 21 | or not math.isfinite(value) |
| 22 | ): |
| 23 | raise ValueError("sampled token logprob must be finite") |
| 24 | object.__setattr__(self, "sampled_token_logprob", float(value)) |
| 25 | |
| 26 | |
| 27 | def oracles_match(expected: Oracle, actual: Oracle) -> bool: |
| 28 | return ( |
| 29 | expected.token_ids == actual.token_ids |
| 30 | and expected.text == actual.text |
| 31 | and math.isclose( |
| 32 | expected.sampled_token_logprob, |
| 33 | actual.sampled_token_logprob, |
| 34 | rel_tol=0.0, |
| 35 | abs_tol=1e-3, |
| 36 | ) |
| 37 | ) |