Spaces:
Running on Zero
Running on Zero
| """A minimal right-context copy test for bidirectional denoisers.""" | |
| from __future__ import annotations | |
| from typing import Any | |
| import torch | |
| from .inference import _prompt_ids, forward_denoising | |
| def _encode(tokenizer: Any, text: str) -> list[int]: | |
| """Encode one deliberately boundary-aware benchmark segment.""" | |
| return list(tokenizer.encode(text, add_special_tokens=False)) | |
| def _encoded_example(session: Any, example: Any) -> tuple[list[int], list[int], list[int]]: | |
| """Return chat/prefix IDs, clean target IDs, and visible suffix IDs.""" | |
| system_prompt = str(example.metadata.get("system_prompt", "You are a helpful assistant.")) | |
| chat = _prompt_ids(session.tokenizer, example.prompt, system_prompt, session.prompt_format) | |
| answer_prefix = _encode(session.tokenizer, str(example.metadata["answer_prefix"])) | |
| target = _encode(session.tokenizer, str(example.metadata["target_text"])) | |
| suffix = _encode(session.tokenizer, str(example.metadata["answer_suffix"])) | |
| if not target: | |
| raise ValueError(f"Infilling example {example.example_id} has an empty tokenized target") | |
| return chat + answer_prefix, target, suffix | |
| def score_bidirectional_example(session: Any, example: Any) -> dict[str, Any]: | |
| """Predict an earlier masked word with and without the later copy clue.""" | |
| left, target, suffix = _encoded_example(session, example) | |
| target_start = len(left) | |
| masked_target = [int(session.mask_token_id)] * len(target) | |
| def predict(input_ids: list[int]) -> tuple[list[int], int]: | |
| tokens = torch.tensor([input_ids], dtype=torch.long, device=session.device) | |
| padding = torch.zeros_like(tokens, dtype=torch.bool) | |
| logits = forward_denoising(session, tokens, padding)[0, target_start:target_start + len(target)].float() | |
| predicted = logits.argmax(dim=-1).tolist() | |
| correct = sum(int(actual == expected) for actual, expected in zip(predicted, target)) | |
| return predicted, correct | |
| predicted, correct = predict(left + masked_target + suffix) | |
| without_clue_predicted, without_clue_correct = predict(left + masked_target) | |
| return { | |
| "target_ids": target, | |
| "prediction_ids": predicted, | |
| "prediction": session.tokenizer.decode(predicted, skip_special_tokens=True).strip(), | |
| "target": str(example.metadata["target_text"]).strip(), | |
| "target_tokens": len(target), | |
| "correct_tokens": correct, | |
| "exact_match": predicted == target, | |
| "without_clue_prediction_ids": without_clue_predicted, | |
| "without_clue_prediction": session.tokenizer.decode(without_clue_predicted, skip_special_tokens=True).strip(), | |
| "without_clue_correct_tokens": without_clue_correct, | |
| "without_clue_exact_match": without_clue_predicted == target, | |
| "clue_changed_prediction": predicted != without_clue_predicted, | |
| } | |
| def summarize_infilling(results: list[dict[str, Any]]) -> dict[str, Any]: | |
| """Report accuracy with the later clue, without it, and their difference.""" | |
| examples = len(results) | |
| tokens = sum(int(result["target_tokens"]) for result in results) | |
| exact_match = sum(int(result["exact_match"]) for result in results) / max(examples, 1) | |
| token_accuracy = sum(int(result["correct_tokens"]) for result in results) / max(tokens, 1) | |
| without_clue_exact = sum(int(result["without_clue_exact_match"]) for result in results) / max(examples, 1) | |
| without_clue_token = sum(int(result["without_clue_correct_tokens"]) for result in results) / max(tokens, 1) | |
| clue_changed = sum(int(result["clue_changed_prediction"]) for result in results) / max(examples, 1) | |
| summary = { | |
| "accuracy": exact_match, | |
| "exact_match": exact_match, | |
| "token_accuracy": token_accuracy, | |
| "without_clue_exact_match": without_clue_exact, | |
| "without_clue_token_accuracy": without_clue_token, | |
| "right_context_gain_exact_match": exact_match - without_clue_exact, | |
| "right_context_gain_token_accuracy": token_accuracy - without_clue_token, | |
| "clue_changed_prediction_rate": clue_changed, | |
| "correct": sum(int(result["exact_match"]) for result in results), | |
| "total": examples, | |
| "tokens": tokens, | |
| } | |
| return summary | |