from abc import ABC, abstractmethod
from typing import Dict, Optional
from src.etl.types import QAExample, Prediction


class QAModel(ABC):
    """Basic contract dictating specific QA model implementation requirements."""

    @abstractmethod
    def train(
        self,
        train_examples: Dict[str, QAExample],
        val_examples: Optional[Dict[str, QAExample]] = None,
    ) -> None:
        """
        Trains the model; assumes uniqueness of keys of train_examples (unique question IDs).
        """
        raise NotImplementedError

    @abstractmethod
    def predict(self, examples: Dict[str, QAExample]) -> Dict[str, Prediction]:
        """
        Produces one Prediction per question ID.
        """
        raise NotImplementedError