Spaces:
Running
Running
File size: 1,590 Bytes
f4c6857 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """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]"
@dataclass(frozen=True)
class CiteInput:
input_ids: list[int]
block_positions: list[int]
n_blocks: int
@property
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))
|