2022-11-03 09:07:54 -06:00
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
from typing import List, Tuple, Optional, TypeVar, Type
|
2022-10-28 11:24:00 -06:00
|
|
|
|
|
|
|
from text_generation.models.types import Batch, GeneratedText
|
|
|
|
|
2022-11-03 09:07:54 -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):
|
|
|
|
@property
|
|
|
|
@abstractmethod
|
|
|
|
def batch_type(self) -> Type[B]:
|
|
|
|
raise NotImplementedError
|
2022-10-28 11:24:00 -06:00
|
|
|
|
2022-11-03 09:07:54 -06:00
|
|
|
@abstractmethod
|
2022-10-28 11:24:00 -06:00
|
|
|
def generate_token(
|
2022-11-03 09:07:54 -06:00
|
|
|
self, batch: B
|
|
|
|
) -> Tuple[List[GeneratedText], Optional[B]]:
|
|
|
|
raise NotImplementedError
|