Spaces:
Running
Running
| """Cite input format: ``[CLS] question [SEP] [BLOCK] text [BLOCK] text ... [SEP]``.""" | |
| from __future__ import annotations | |
| from collections.abc import Sequence | |
| from dataclasses import dataclass | |
| from typing import Any | |
| BLOCK_TOKEN = "[BLOCK]" | |
| class CiteInput: | |
| input_ids: list[int] | |
| block_positions: list[int] | |
| n_blocks: int | |
| def truncated_blocks(self) -> int: | |
| """Blocks left out because they did not fit in ``max_tokens``.""" | |
| return self.n_blocks - len(self.block_positions) | |
| def add_block_token(tokenizer: Any) -> int: | |
| tokenizer.add_special_tokens({"additional_special_tokens": [BLOCK_TOKEN]}) | |
| return int(tokenizer.convert_tokens_to_ids(BLOCK_TOKEN)) | |
| def encode( | |
| tokenizer: Any, | |
| block_token_id: int, | |
| question: str, | |
| blocks: Sequence[str], | |
| *, | |
| max_tokens: int, | |
| ) -> CiteInput: | |
| """Encode whole blocks in order until the next one would not fit.""" | |
| ids = [tokenizer.cls_token_id, *tokenizer.encode(question, add_special_tokens=False)] | |
| ids.append(tokenizer.sep_token_id) | |
| if len(ids) + 1 > max_tokens: | |
| raise ValueError(f"question alone needs {len(ids) + 1} tokens, over {max_tokens}") | |
| positions = [] | |
| for text in blocks: | |
| piece = [block_token_id, *tokenizer.encode(text, add_special_tokens=False)] | |
| if len(ids) + len(piece) + 1 > max_tokens: | |
| break | |
| positions.append(len(ids)) | |
| ids.extend(piece) | |
| ids.append(tokenizer.sep_token_id) | |
| return CiteInput(input_ids=ids, block_positions=positions, n_blocks=len(blocks)) | |