models/base_model.py

"""Base model class for Myanmar Ghost project."""

import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


class BaseModel(ABC, nn.Module):
    """Abstract base class for all models."""
    
    def __init__(self, config: Optional[Dict] = None):
        super().__init__()
        self.config = config or {}
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
    
    @abstractmethod
    def forward(self, *args, **kwargs) -> torch.Tensor:
        """Forward pass."""
        pass
    
    @abstractmethod
    def predict(self, *args, **kwargs) -> Dict[str, Any]:
        """Make predictions."""
        pass
    
    def save(self, path: str) -> None:
        """Save model checkpoint."""
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        torch.save({
            "model_state_dict": self.state_dict(),
            "config": self.config,
        }, path)
        logger.info(f"Model saved to {path}")
    
    def load(self, path: str) -> None:
        """Load model checkpoint."""
        checkpoint = torch.load(path, map_location=self.device)
        self.load_state_dict(checkpoint["model_state_dict"])
        if "config" in checkpoint:
            self.config = checkpoint["config"]
        logger.info(f"Model loaded from {path}")
    
    def get_num_parameters(self) -> int:
        """Get total number of parameters."""
        return sum(p.numel() for p in self.parameters())
    
    def get_num_trainable_parameters(self) -> int:
        """Get number of trainable parameters."""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


class SentimentClassifier(nn.Module):
    """Base sentiment classifier."""
    
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        num_classes: int = 4,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim // 2)
        self.fc3 = nn.Linear(hidden_dim // 2, num_classes)
        self.relu = nn.ReLU()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x


def create_model(
    model_type: str = "transformer",
    **kwargs,
) -> BaseModel:
    """Factory function to create models."""
    from .transformer_model import TransformerSentimentModel
    from .multimodal_model import MultiModalSentimentModel
    
    if model_type == "transformer":
        return TransformerSentimentModel(**kwargs)
    elif model_type == "multimodal":
        return MultiModalSentimentModel(**kwargs)
    elif model_type == "base":
        return SentimentClassifier(**kwargs)
    else:
        raise ValueError(f"Unknown model type: {model_type}")


if __name__ == "__main__":
    model = SentimentClassifier(input_dim=768, hidden_dim=256, num_classes=4)
    print(f"Model parameters: {model.get_num_parameters():,}")