src/models/base.py

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional
import gradio as gr


class BaseModel(ABC):
    """Base class for all ML/AI models."""
    
    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description
        self._model: Optional[Any] = None
        self._interface: Optional[gr.Interface] = None
    
    @abstractmethod
    def load_model(self) -> None:
        """Load the model into memory."""
        pass
    
    @abstractmethod
    def create_interface(self) -> gr.Interface:
        """Create and return a Gradio interface for the model."""
        pass
    
    @abstractmethod
    async def predict(self, *args, **kwargs) -> Any:
        """Make predictions using the model."""
        pass
    
    def get_interface(self) -> gr.Interface:
        """Get or create the Gradio interface."""
        if self._interface is None:
            self._interface = self.create_interface()
        return self._interface
    
    def get_info(self) -> Dict[str, str]:
        """Get model information."""
        return {
            "name": self.name,
            "description": self.description,
            "status": "loaded" if self._model is not None else "unloaded"
        }