model.py

import torch

from abc import ABC, abstractmethod
from typing import List, Tuple, Optional, TypeVar, Type
from tokenizers import Tokenizer

from text_generation.models.types import Batch, GeneratedText

B = TypeVar("B", bound=Batch)


class Model(ABC):
    def __init__(self, tokenizer: Tokenizer, num_heads: int, device: torch.device):
        self.tokenizer = tokenizer
        self.num_heads = num_heads
        self.device = device

    @property
    @abstractmethod
    def batch_type(self) -> Type[B]:
        raise NotImplementedError

    @abstractmethod
    def generate_token(self, batch: B) -> Tuple[List[GeneratedText], Optional[B]]:
        raise NotImplementedError