Spaces:
Running on Zero
Running on Zero
| """Verify whole-sequence decoding independently of any private checkpoint.""" | |
| import unittest | |
| from types import SimpleNamespace | |
| import torch | |
| from tokenizers import Tokenizer | |
| from tokenizers.decoders import ByteLevel | |
| from tokenizers.models import BPE | |
| from transformers import PreTrainedTokenizerFast | |
| from vendor.attention_extractor import extract_generation_attention | |
| class ScriptedModel: | |
| """A test double emitting byte fragments and an EOS with uniform attention.""" | |
| device = torch.device("cpu") | |
| def __init__(self, ids: list[int]) -> None: | |
| """Provide deterministic next tokens without a model download.""" | |
| self.ids = ids | |
| self.calls = 0 | |
| def __call__(self, input_ids: torch.Tensor, output_attentions: bool = True) -> SimpleNamespace: | |
| """Return one greedy choice and an actual tensor for the extraction loop.""" | |
| size = input_ids.shape[1] | |
| logits = torch.zeros(1, size, 4) | |
| logits[0, -1, self.ids[self.calls]] = 1 | |
| self.calls += 1 | |
| return SimpleNamespace(logits=logits, attentions=[torch.full((1, 1, size, size), 1 / size)]) | |
| class GenerationTextTests(unittest.TestCase): | |
| """Check Unicode, special tokens and stopping without changing attention.""" | |
| def setUp(self) -> None: | |
| """Use a real byte-level decoder: two pieces are needed for one character.""" | |
| tokenizer = Tokenizer(BPE({"Ã": 0, "©": 1, "[EOS]": 2, "x": 3}, merges=[])) | |
| tokenizer.decoder = ByteLevel() | |
| self.tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer, eos_token="[EOS]") | |
| def extract(self, ids: list[int], limit: int) -> dict: | |
| """Exercise the production extractor, not an alternate decoding path.""" | |
| model = ScriptedModel(ids) | |
| result = extract_generation_attention( | |
| model, | |
| self.tokenizer, | |
| "x", | |
| [], | |
| ["instr_t"], | |
| {"layer_types": ["attention"], "num_heads": 1}, | |
| max_tokens=limit, | |
| layer=0, | |
| ) | |
| self.assertEqual(model.calls, len(ids), "Readable output must not trigger additional model calls") | |
| return result | |
| def test_unicode_sequence_and_eos(self) -> None: | |
| """Readable text must decode IDs together, while the raw strip keeps EOS.""" | |
| result = self.extract([0, 1, 2], 64) | |
| self.assertEqual(result["generated_text"], "é") | |
| self.assertNotEqual("".join(result["generated_tokens"][:-1]), "é") | |
| self.assertEqual(result["generated_tokens"][-1], "[EOS]") | |
| self.assertEqual(result["generated_token_ids"], [0, 1, 2]) | |
| self.assertEqual(result["generation_stop_reason"], "eos") | |
| self.assertEqual(len(result["attention_per_step"]), 3) | |
| self.assertEqual(result["attention_per_step"][0]["span_attention"], {"instr_t": 1.0, "generated": 0.0}) | |
| self.assertEqual(result["attention_per_step"][1]["span_attention"], {"instr_t": 0.5, "generated": 0.5}) | |
| def test_token_limit_does_not_claim_completion(self) -> None: | |
| """The limit is distinct from the model choosing to end its response.""" | |
| result = self.extract([0, 1], 2) | |
| self.assertEqual(result["generated_text"], "é") | |
| self.assertEqual(result["generation_stop_reason"], "token_limit") | |
| def test_empty_readable_output_stays_empty(self) -> None: | |
| """An EOS-only response must never be replaced with an expected answer.""" | |
| result = self.extract([2], 64) | |
| self.assertEqual(result["generated_text"], "") | |
| self.assertEqual(result["generated_tokens"], ["[EOS]"]) | |
| self.assertEqual(result["generation_stop_reason"], "eos") | |
| if __name__ == "__main__": | |
| unittest.main() | |