core/base.py

from typing import List, Optional


class RerankerModel:
    """Base class for reranking models."""
    
    def __init__(self, model_id: str, model_name: str, model_type: str):
        self.model_id = model_id
        self.model_name = model_name
        self.model_type = model_type
        self.model = None
        self.tokenizer = None
        self.loaded = False
    
    def load(self):
        """Load the model. To be implemented by subclasses."""
        raise NotImplementedError
    
    def rerank(self, query: str, documents: List[str], instruction: Optional[str] = None) -> List[float]:
        """Rerank documents. To be implemented by subclasses."""
        raise NotImplementedError