File size: 3,267 Bytes
cfb5e7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """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():,}")
|