2022-11-04 07:22:47 -06:00
|
|
|
import torch
|
|
|
|
|
2022-11-03 09:07:54 -06:00
|
|
|
from abc import ABC, abstractmethod
|
2022-11-04 11:03:04 -06:00
|
|
|
from typing import List, Tuple, Optional, TypeVar, Type
|
2023-01-17 01:10:22 -07:00
|
|
|
from transformers import PreTrainedTokenizerBase
|
2022-10-28 11:24:00 -06:00
|
|
|
|
|
|
|
from text_generation.models.types import Batch, GeneratedText
|
|
|
|
|
2022-11-04 11:03:04 -06:00
|
|
|
B = TypeVar("B", bound=Batch)
|
|
|
|
|
2022-10-28 11:24:00 -06:00
|
|
|
|
2022-11-03 09:07:54 -06:00
|
|
|
class Model(ABC):
|
2023-01-17 01:10:22 -07:00
|
|
|
def __init__(self, tokenizer: PreTrainedTokenizerBase, device: torch.device):
|
2022-11-04 07:22:47 -06:00
|
|
|
self.tokenizer = tokenizer
|
2023-02-24 07:55:57 -07:00
|
|
|
self.all_special_ids = set(tokenizer.all_special_ids)
|
2022-11-04 07:22:47 -06:00
|
|
|
self.device = device
|
|
|
|
|
2023-03-06 05:22:58 -07:00
|
|
|
# see `decode_token` method
|
|
|
|
self.tokenizer.add_special_tokens(
|
|
|
|
{"additional_special_tokens": ["<decode-token>"]}
|
|
|
|
)
|
|
|
|
self.special_decode_token_id = self.tokenizer.convert_tokens_to_ids(
|
|
|
|
"<decode-token>"
|
|
|
|
)
|
|
|
|
self.special_decode_token_length = len("<decode-token>")
|
|
|
|
|
2022-11-04 11:03:04 -06:00
|
|
|
@property
|
2022-11-03 09:07:54 -06:00
|
|
|
@abstractmethod
|
2022-11-04 11:03:04 -06:00
|
|
|
def batch_type(self) -> Type[B]:
|
2022-11-03 09:07:54 -06:00
|
|
|
raise NotImplementedError
|
2022-10-28 11:24:00 -06:00
|
|
|
|
2022-11-04 11:03:04 -06:00
|
|
|
@abstractmethod
|
|
|
|
def generate_token(self, batch: B) -> Tuple[List[GeneratedText], Optional[B]]:
|
|
|
|
raise NotImplementedError
|
2023-03-06 05:22:58 -07:00
|
|
|
|
|
|
|
def decode_token(self, token_id: int) -> str:
|
|
|
|
"""Hack to hopefully support generate_stream for the maximum number of tokenizers"""
|
|
|
|
# append token to special decode token and decode both
|
|
|
|
result = self.tokenizer.decode(
|
|
|
|
[self.special_decode_token_id, token_id], skip_special_tokens=False
|
|
|
|
)
|
|
|
|
# slice to remove special decode token
|
|
|
|
return result[self.special_decode_token_length :]
|