Upload models.py with huggingface_hub

models.py

@@ -0,0 +1,978 @@
+"""
+Model architectures cho Aspect-Based Sentiment Analysis
+Hỗ trợ nhiều architectures: Transformer-based, CNN, LSTM, và hybrid models
+"""
+import torch
+import os
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import (
+    RobertaPreTrainedModel, RobertaModel, 
+    BertPreTrainedModel, BertModel,
+    XLMRobertaPreTrainedModel, XLMRobertaModel,
+    BartPreTrainedModel, BartModel, BartForSequenceClassification,
+    T5PreTrainedModel, T5EncoderModel,
+    AutoConfig, AutoModel, AutoTokenizer,
+    PreTrainedModel
+)
+from transformers.modeling_outputs import SequenceClassifierOutput
+from typing import Optional
+
+
+class BaseABSA(PreTrainedModel):
+    """Base class cho tất cả ABSA models"""
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_aspects = config.num_aspects
+        self.num_sentiments = config.num_sentiments
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None):
+        raise NotImplementedError
+
+    def get_sentiment_classifiers(self, hidden_size):
+        """Create sentiment classifiers cho từng aspect"""
+        return nn.ModuleList([
+            nn.Linear(hidden_size, self.num_sentiments + 1)  # +1 cho "none"
+            for _ in range(self.num_aspects)
+        ])
+
+
+# ========== Transformer-based Models ==========
+
+class TransformerForABSA(RobertaPreTrainedModel):
+    """RoBERTa-based model (cho PhoBERT, ViSoBERT, RoBERTa-GRU)"""
+    base_model_prefix = "roberta"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = RobertaModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(config.hidden_size, config.num_sentiments + 1)
+            for _ in range(config.num_aspects)
+        ])
+        self.init_weights()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        # RoBERTa-based models don't use token_type_ids, so we ignore it
+        kwargs.pop('token_type_ids', None)
+        # Filter kwargs to only include valid arguments for RobertaModel
+        model_kwargs = {
+            k: v for k, v in kwargs.items() 
+            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
+                    'output_attentions', 'output_hidden_states']
+        }
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.roberta(input_ids, attention_mask=attention_mask, return_dict=return_dict, **model_kwargs)
+        pooled = self.dropout(outputs.pooler_output)
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            B, A, _ = all_logits.size()
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits_flat, targets_flat)
+
+        if not return_dict:
+            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
+        
+        # T5 returns BaseModelOutput, which has hidden_states
+        # But we need to handle it properly
+        hidden_states = getattr(outputs, 'hidden_states', None)
+        attentions = getattr(outputs, 'attentions', None)
+        
+        return SequenceClassifierOutput(
+            loss=loss, logits=all_logits,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        # Ensure directory exists
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save backbone
+        self.roberta.save_pretrained(save_directory, **kwargs)
+        
+        # Update and save config with custom attributes
+        config = self.roberta.config
+        config.num_aspects = len(self.sentiment_classifiers)
+        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1  # -1 vì không tính lớp "none"
+        # Auto map để AutoModel tự động load đúng class
+        # models.py sẽ được upload vào root của repo
+        config.auto_map = {
+            "AutoModel": "models.TransformerForABSA",
+            "AutoModelForSequenceClassification": "models.TransformerForABSA"
+        }
+        # Lưu thêm thông tin vào config để dễ dàng load lại
+        if not hasattr(config, 'custom_model_type'):
+            config.custom_model_type = 'TransformerForABSA'
+        config.save_pretrained(save_directory, **kwargs)
+        
+        # Save full state_dict (bao gồm cả sentiment_classifiers)
+        sd = kwargs.get("state_dict", None) or self.state_dict()
+        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        
+        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
+        if num_aspects is None:
+            num_aspects = getattr(config, 'num_aspects', None)
+            if num_aspects is None:
+                raise ValueError("num_aspects must be provided or present in config")
+        
+        if num_sentiments is None:
+            num_sentiments = getattr(config, 'num_sentiments', None)
+            if num_sentiments is None:
+                raise ValueError("num_sentiments must be provided or present in config")
+        
+        config.num_aspects = num_aspects
+        config.num_sentiments = num_sentiments
+        
+        model = cls(config)
+        
+        # Load backbone weights
+        model.roberta = RobertaModel.from_pretrained(
+            pretrained_model_name_or_path, config=config,
+            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
+        )
+        
+        # Load full state_dict nếu có (bao gồm sentiment_classifiers)
+        try:
+            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+            if os.path.exists(state_dict_path):
+                state_dict = torch.load(state_dict_path, map_location="cpu")
+                model.load_state_dict(state_dict, strict=False)
+            elif "state_dict" in kwargs:
+                model.load_state_dict(kwargs["state_dict"], strict=False)
+        except Exception as e:
+            print(f"⚠ Warning: Could not load full state_dict: {e}")
+        
+        return model
+
+
+class BERTForABSA(BertPreTrainedModel):
+    """BERT-based model (cho mBERT)"""
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = BertModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(config.hidden_size, config.num_sentiments + 1)
+            for _ in range(config.num_aspects)
+        ])
+        self.init_weights()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, token_type_ids=None, **kwargs):
+        # BERT models can use token_type_ids, but for single sentence tasks, it's usually all zeros
+        # Filter kwargs to only include valid arguments for BertModel
+        model_kwargs = {
+            k: v for k, v in kwargs.items() 
+            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
+                    'output_attentions', 'output_hidden_states']
+        }
+        # BERT expects token_type_ids, but if not provided, it will default to all zeros
+        if token_type_ids is not None:
+            model_kwargs['token_type_ids'] = token_type_ids
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.bert(input_ids, attention_mask=attention_mask, return_dict=return_dict, **model_kwargs)
+        pooled = self.dropout(outputs.pooler_output)
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
+
+        if not return_dict:
+            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
+        
+        # T5 returns BaseModelOutput, which has hidden_states
+        # But we need to handle it properly
+        hidden_states = getattr(outputs, 'hidden_states', None)
+        attentions = getattr(outputs, 'attentions', None)
+        
+        return SequenceClassifierOutput(
+            loss=loss, logits=all_logits,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        """Save model with custom attributes"""
+        os.makedirs(save_directory, exist_ok=True)
+        self.bert.save_pretrained(save_directory, **kwargs)
+        config = self.bert.config
+        config.num_aspects = len(self.sentiment_classifiers)
+        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1
+        config.auto_map = {
+            "AutoModel": "models.BERTForABSA",
+            "AutoModelForSequenceClassification": "models.BERTForABSA"
+        }
+        if not hasattr(config, 'custom_model_type'):
+            config.custom_model_type = 'BERTForABSA'
+        config.save_pretrained(save_directory, **kwargs)
+        sd = kwargs.get("state_dict", None) or self.state_dict()
+        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        
+        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
+        if num_aspects is None:
+            num_aspects = getattr(config, 'num_aspects', None)
+            if num_aspects is None:
+                raise ValueError("num_aspects must be provided or present in config")
+        
+        if num_sentiments is None:
+            num_sentiments = getattr(config, 'num_sentiments', None)
+            if num_sentiments is None:
+                raise ValueError("num_sentiments must be provided or present in config")
+        
+        config.num_aspects = num_aspects
+        config.num_sentiments = num_sentiments
+        model = cls(config)
+        model.bert = BertModel.from_pretrained(
+            pretrained_model_name_or_path, config=config,
+            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
+        )
+        
+        # Load full state_dict if available
+        try:
+            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+            if os.path.exists(state_dict_path):
+                state_dict = torch.load(state_dict_path, map_location="cpu")
+                model.load_state_dict(state_dict, strict=False)
+            elif "state_dict" in kwargs:
+                model.load_state_dict(kwargs["state_dict"], strict=False)
+        except Exception as e:
+            print(f"⚠ Warning: Could not load full state_dict: {e}")
+        
+        return model
+
+
+class XLMRobertaForABSA(XLMRobertaPreTrainedModel):
+    """XLM-RoBERTa-based model"""
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = XLMRobertaModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(config.hidden_size, config.num_sentiments + 1)
+            for _ in range(config.num_aspects)
+        ])
+        self.init_weights()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.roberta(input_ids, attention_mask=attention_mask, return_dict=return_dict)
+        pooled = self.dropout(outputs.pooler_output)
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
+
+        if not return_dict:
+            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
+        
+        # T5 returns BaseModelOutput, which has hidden_states
+        # But we need to handle it properly
+        hidden_states = getattr(outputs, 'hidden_states', None)
+        attentions = getattr(outputs, 'attentions', None)
+        
+        return SequenceClassifierOutput(
+            loss=loss, logits=all_logits,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        
+        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
+        if num_aspects is None:
+            num_aspects = getattr(config, 'num_aspects', None)
+            if num_aspects is None:
+                raise ValueError("num_aspects must be provided or present in config")
+        
+        if num_sentiments is None:
+            num_sentiments = getattr(config, 'num_sentiments', None)
+            if num_sentiments is None:
+                raise ValueError("num_sentiments must be provided or present in config")
+        
+        config.num_aspects = num_aspects
+        config.num_sentiments = num_sentiments
+        model = cls(config)
+        model.roberta = XLMRobertaModel.from_pretrained(
+            pretrained_model_name_or_path, config=config,
+            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
+        )
+        
+        # Load full state_dict nếu có (bao gồm sentiment_classifiers)
+        try:
+            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+            if os.path.exists(state_dict_path):
+                state_dict = torch.load(state_dict_path, map_location="cpu")
+                model.load_state_dict(state_dict, strict=False)
+            elif "state_dict" in kwargs:
+                model.load_state_dict(kwargs["state_dict"], strict=False)
+        except Exception as e:
+            print(f"⚠ Warning: Could not load full state_dict: {e}")
+        
+        return model
+
+
+class RoBERTaGRUForABSA(RobertaPreTrainedModel):
+    """RoBERTa + GRU hybrid model"""
+    base_model_prefix = "roberta"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = RobertaModel(config)
+        self.gru = nn.GRU(
+            config.hidden_size, config.hidden_size, 
+            num_layers=2, batch_first=True, bidirectional=True, dropout=0.2
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(config.hidden_size * 2, config.num_sentiments + 1)  # *2 vì bidirectional
+            for _ in range(config.num_aspects)
+        ])
+        self.init_weights()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.roberta(input_ids, attention_mask=attention_mask, return_dict=return_dict)
+        
+        # Use last_hidden_state thay vì pooler_output
+        sequence_output = outputs.last_hidden_state  # [B, L, H]
+        
+        # GRU layer
+        gru_out, _ = self.gru(sequence_output)  # [B, L, H*2]
+        # Take last timestep
+        pooled = self.dropout(gru_out[:, -1, :])  # [B, H*2]
+        
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
+
+        if not return_dict:
+            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
+        
+        # T5 returns BaseModelOutput, which has hidden_states
+        # But we need to handle it properly
+        hidden_states = getattr(outputs, 'hidden_states', None)
+        attentions = getattr(outputs, 'attentions', None)
+        
+        return SequenceClassifierOutput(
+            loss=loss, logits=all_logits,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        
+        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
+        if num_aspects is None:
+            num_aspects = getattr(config, 'num_aspects', None)
+            if num_aspects is None:
+                raise ValueError("num_aspects must be provided or present in config")
+        
+        if num_sentiments is None:
+            num_sentiments = getattr(config, 'num_sentiments', None)
+            if num_sentiments is None:
+                raise ValueError("num_sentiments must be provided or present in config")
+        
+        config.num_aspects = num_aspects
+        config.num_sentiments = num_sentiments
+        model = cls(config)
+        model.roberta = RobertaModel.from_pretrained(
+            pretrained_model_name_or_path, config=config,
+            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
+        )
+        
+        # Load full state_dict nếu có (bao gồm GRU và sentiment_classifiers)
+        try:
+            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+            if os.path.exists(state_dict_path):
+                state_dict = torch.load(state_dict_path, map_location="cpu")
+                model.load_state_dict(state_dict, strict=False)
+            elif "state_dict" in kwargs:
+                model.load_state_dict(kwargs["state_dict"], strict=False)
+        except Exception as e:
+            print(f"⚠ Warning: Could not load full state_dict: {e}")
+        
+        return model
+
+
+class BartForABSA(BartPreTrainedModel):
+    """BART-based model (cho BartPho)"""
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = BartModel(config)
+        self.dropout = nn.Dropout(config.dropout)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(config.d_model, config.num_sentiments + 1)
+            for _ in range(config.num_aspects)
+        ])
+        self.init_weights()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        # BART models don't use token_type_ids, so we ignore it
+        kwargs.pop('token_type_ids', None)
+        # Filter kwargs to only include valid arguments for BartModel
+        # Remove training-specific arguments that BartModel doesn't accept
+        model_kwargs = {
+            k: v for k, v in kwargs.items() 
+            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
+                    'output_attentions', 'output_hidden_states']
+        }
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # IMPORTANT: For BART, we need to access encoder output directly
+        # BartModel.forward() returns decoder output in last_hidden_state
+        # We need to call encoder separately to get encoder hidden states
+        # Only call encoder once (don't call full model.forward() to avoid double computation)
+        encoder_outputs = self.model.get_encoder()(
+            input_ids, 
+            attention_mask=attention_mask, 
+            return_dict=True,
+            **{k: v for k, v in model_kwargs.items()}
+        )
+        sequence_output = encoder_outputs.last_hidden_state  # [B, L, H] - encoder output
+        
+        # Mean pooling with attention mask (weighted mean to avoid padding tokens)
+        if attention_mask is not None:
+            # Expand attention mask to match sequence_output dimensions
+            attention_mask_expanded = attention_mask.unsqueeze(-1).expand(sequence_output.size()).float()
+            # Sum over sequence length, divide by number of non-padding tokens
+            sum_embeddings = torch.sum(sequence_output * attention_mask_expanded, dim=1)
+            sum_mask = torch.clamp(attention_mask_expanded.sum(dim=1), min=1e-9)
+            pooled = sum_embeddings / sum_mask  # [B, H]
+        else:
+            pooled = sequence_output.mean(dim=1)  # [B, H]
+        
+        pooled = self.dropout(pooled)
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
+
+        if not return_dict:
+            return ((loss, all_logits) + (encoder_outputs.hidden_states, encoder_outputs.attentions)) if loss is not None else (all_logits,)
+        
+        # Use encoder outputs for hidden_states and attentions
+        hidden_states = getattr(encoder_outputs, 'hidden_states', None)
+        attentions = getattr(encoder_outputs, 'attentions', None)
+        
+        return SequenceClassifierOutput(
+            loss=loss, logits=all_logits,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        """Save model with custom attributes"""
+        os.makedirs(save_directory, exist_ok=True)
+        self.model.save_pretrained(save_directory, **kwargs)
+        config = self.model.config
+        config.num_aspects = len(self.sentiment_classifiers)
+        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1
+        config.auto_map = {
+            "AutoModel": "models.BartForABSA",
+            "AutoModelForSequenceClassification": "models.BartForABSA"
+        }
+        if not hasattr(config, 'custom_model_type'):
+            config.custom_model_type = 'BartForABSA'
+        config.save_pretrained(save_directory, **kwargs)
+        sd = kwargs.get("state_dict", None) or self.state_dict()
+        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        
+        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
+        if num_aspects is None:
+            num_aspects = getattr(config, 'num_aspects', None)
+            if num_aspects is None:
+                raise ValueError("num_aspects must be provided or present in config")
+        
+        if num_sentiments is None:
+            num_sentiments = getattr(config, 'num_sentiments', None)
+            if num_sentiments is None:
+                raise ValueError("num_sentiments must be provided or present in config")
+        
+        config.num_aspects = num_aspects
+        config.num_sentiments = num_sentiments
+        model = cls(config)
+        model.model = BartModel.from_pretrained(
+            pretrained_model_name_or_path, config=config,
+            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
+        )
+        
+        # Load full state_dict if available
+        try:
+            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+            if os.path.exists(state_dict_path):
+                state_dict = torch.load(state_dict_path, map_location="cpu")
+                model.load_state_dict(state_dict, strict=False)
+            elif "state_dict" in kwargs:
+                model.load_state_dict(kwargs["state_dict"], strict=False)
+        except Exception as e:
+            print(f"⚠ Warning: Could not load full state_dict: {e}")
+        
+        return model
+
+
+class T5ForABSA(T5PreTrainedModel):
+    """T5-based model (cho ViT5) - sử dụng encoder only"""
+    def __init__(self, config):
+        super().__init__(config)
+        self.encoder = T5EncoderModel(config)
+        self.dropout = nn.Dropout(config.dropout_rate)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(config.d_model, config.num_sentiments + 1)
+            for _ in range(config.num_aspects)
+        ])
+        self.init_weights()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        # T5 models don't use token_type_ids, so we ignore it
+        kwargs.pop('token_type_ids', None)
+        # Filter kwargs to only include valid arguments for T5EncoderModel
+        model_kwargs = {
+            k: v for k, v in kwargs.items() 
+            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
+                    'output_attentions', 'output_hidden_states']
+        }
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.encoder(input_ids, attention_mask=attention_mask, return_dict=return_dict, **model_kwargs)
+        
+        # Mean pooling with attention mask (weighted mean to avoid padding tokens)
+        sequence_output = outputs.last_hidden_state  # [B, L, H]
+        if attention_mask is not None:
+            # Expand attention mask to match sequence_output dimensions
+            attention_mask_expanded = attention_mask.unsqueeze(-1).expand(sequence_output.size()).float()
+            # Sum over sequence length, divide by number of non-padding tokens
+            sum_embeddings = torch.sum(sequence_output * attention_mask_expanded, dim=1)
+            sum_mask = torch.clamp(attention_mask_expanded.sum(dim=1), min=1e-9)
+            pooled = sum_embeddings / sum_mask  # [B, H]
+        else:
+            pooled = sequence_output.mean(dim=1)  # [B, H]
+        
+        pooled = self.dropout(pooled)
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
+
+        if not return_dict:
+            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
+        
+        # T5 returns BaseModelOutput, which has hidden_states
+        # But we need to handle it properly
+        hidden_states = getattr(outputs, 'hidden_states', None)
+        attentions = getattr(outputs, 'attentions', None)
+        
+        return SequenceClassifierOutput(
+            loss=loss, logits=all_logits,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        """Save model with custom attributes"""
+        os.makedirs(save_directory, exist_ok=True)
+        self.encoder.save_pretrained(save_directory, **kwargs)
+        config = self.encoder.config
+        config.num_aspects = len(self.sentiment_classifiers)
+        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1
+        config.auto_map = {
+            "AutoModel": "models.T5ForABSA",
+            "AutoModelForSequenceClassification": "models.T5ForABSA"
+        }
+        if not hasattr(config, 'custom_model_type'):
+            config.custom_model_type = 'T5ForABSA'
+        config.save_pretrained(save_directory, **kwargs)
+        sd = kwargs.get("state_dict", None) or self.state_dict()
+        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        
+        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
+        if num_aspects is None:
+            num_aspects = getattr(config, 'num_aspects', None)
+            if num_aspects is None:
+                raise ValueError("num_aspects must be provided or present in config")
+        
+        if num_sentiments is None:
+            num_sentiments = getattr(config, 'num_sentiments', None)
+            if num_sentiments is None:
+                raise ValueError("num_sentiments must be provided or present in config")
+        
+        config.num_aspects = num_aspects
+        config.num_sentiments = num_sentiments
+        model = cls(config)
+        model.encoder = T5EncoderModel.from_pretrained(
+            pretrained_model_name_or_path, config=config,
+            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
+        )
+        
+        # Load full state_dict if available
+        try:
+            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+            if os.path.exists(state_dict_path):
+                state_dict = torch.load(state_dict_path, map_location="cpu")
+                model.load_state_dict(state_dict, strict=False)
+            elif "state_dict" in kwargs:
+                model.load_state_dict(kwargs["state_dict"], strict=False)
+        except Exception as e:
+            print(f"⚠ Warning: Could not load full state_dict: {e}")
+        
+        return model
+
+
+# ========== Non-Transformer Models ==========
+
+class TextCNNForABSA(nn.Module):
+    """TextCNN model - không dùng transformers"""
+    def __init__(self, vocab_size, embed_dim, num_filters, filter_sizes, num_aspects, num_sentiments, max_length=256):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.convs = nn.ModuleList([
+            nn.Conv1d(embed_dim, num_filters, kernel_size=fs)
+            for fs in filter_sizes
+        ])
+        self.dropout = nn.Dropout(0.5)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(len(filter_sizes) * num_filters, num_sentiments + 1)
+            for _ in range(num_aspects)
+        ])
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True):
+        # input_ids: [B, L]
+        x = self.embedding(input_ids)  # [B, L, E]
+        x = x.permute(0, 2, 1)  # [B, E, L]
+        
+        conv_outputs = []
+        for conv in self.convs:
+            conv_out = F.relu(conv(x))  # [B, F, L']
+            pooled = F.max_pool1d(conv_out, kernel_size=conv_out.size(2))  # [B, F, 1]
+            conv_outputs.append(pooled.squeeze(2))  # [B, F]
+        
+        x = torch.cat(conv_outputs, dim=1)  # [B, F*len(filter_sizes)]
+        x = self.dropout(x)
+        
+        all_logits = torch.stack([cls(x) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
+
+        if return_dict:
+            return SequenceClassifierOutput(
+                loss=loss, logits=all_logits,
+                hidden_states=None, attentions=None
+            )
+        return (loss, all_logits) if loss is not None else (all_logits,)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        """Load TextCNN model from pretrained path"""
+        import json
+        
+        # Load config
+        config_path = os.path.join(pretrained_model_name_or_path, 'config.json')
+        if os.path.exists(config_path):
+            with open(config_path, 'r', encoding='utf-8') as f:
+                config = json.load(f)
+        else:
+            # Fallback to model_config.json
+            config_path = os.path.join(pretrained_model_name_or_path, 'model_config.json')
+            if os.path.exists(config_path):
+                with open(config_path, 'r', encoding='utf-8') as f:
+                    config = json.load(f)
+            else:
+                raise ValueError(f"Config file not found in {pretrained_model_name_or_path}")
+        
+        # Create model with config
+        model = cls(
+            vocab_size=config.get('vocab_size', 30000),
+            embed_dim=kwargs.get('embed_dim', 300),
+            num_filters=kwargs.get('num_filters', 100),
+            filter_sizes=kwargs.get('filter_sizes', [3, 4, 5]),
+            num_aspects=config['num_aspects'],
+            num_sentiments=config['num_sentiments'],
+            max_length=kwargs.get('max_length', 256)
+        )
+        
+        # Load weights
+        weights_path = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin')
+        if os.path.exists(weights_path):
+            state_dict = torch.load(weights_path, map_location='cpu')
+            model.load_state_dict(state_dict)
+        
+        return model
+
+
+class BiLSTMForABSA(nn.Module):
+    """BiLSTM model - không dùng transformers"""
+    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, num_aspects, num_sentiments, dropout=0.3):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.lstm = nn.LSTM(
+            embed_dim, hidden_dim, num_layers,
+            batch_first=True, bidirectional=True, dropout=dropout
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(hidden_dim * 2, num_sentiments + 1)  # *2 vì bidirectional
+            for _ in range(num_aspects)
+        ])
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True):
+        x = self.embedding(input_ids)  # [B, L, E]
+        lstm_out, (h_n, c_n) = self.lstm(x)  # [B, L, H*2]
+        
+        # Use last non-padding hidden state instead of always using last timestep
+        # This is important because padding tokens can be at the end
+        if attention_mask is not None:
+            # Find the last non-padding token for each sequence
+            # attention_mask: [B, L] where 1 = real token, 0 = padding
+            seq_lengths = attention_mask.sum(dim=1) - 1  # -1 for 0-indexing
+            # Ensure seq_lengths are within valid range
+            seq_lengths = torch.clamp(seq_lengths, min=0, max=lstm_out.size(1) - 1)
+            # Get last hidden state for each sequence: [B, H*2]
+            batch_size = lstm_out.size(0)
+            pooled = lstm_out[torch.arange(batch_size, device=lstm_out.device), seq_lengths, :]
+        else:
+            # Fallback: use last timestep if no attention mask
+            pooled = lstm_out[:, -1, :]  # [B, H*2]
+        
+        pooled = self.dropout(pooled)
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)
+
+        loss = None
+        if labels is not None:
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)
+
+        if return_dict:
+            return SequenceClassifierOutput(
+                loss=loss, logits=all_logits,
+                hidden_states=None, attentions=None
+            )
+        return (loss, all_logits) if loss is not None else (all_logits,)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        """Load BiLSTM model from pretrained path"""
+        import json
+        
+        # Load config
+        config_path = os.path.join(pretrained_model_name_or_path, 'config.json')
+        if os.path.exists(config_path):
+            with open(config_path, 'r', encoding='utf-8') as f:
+                config = json.load(f)
+        else:
+            # Fallback to model_config.json
+            config_path = os.path.join(pretrained_model_name_or_path, 'model_config.json')
+            if os.path.exists(config_path):
+                with open(config_path, 'r', encoding='utf-8') as f:
+                    config = json.load(f)
+            else:
+                raise ValueError(f"Config file not found in {pretrained_model_name_or_path}")
+        
+        # Create model with config
+        model = cls(
+            vocab_size=config.get('vocab_size', 30000),
+            embed_dim=kwargs.get('embed_dim', 300),
+            hidden_dim=kwargs.get('hidden_dim', 256),
+            num_layers=kwargs.get('num_layers', 2),
+            num_aspects=config['num_aspects'],
+            num_sentiments=config['num_sentiments'],
+            dropout=kwargs.get('dropout', 0.3)
+        )
+        
+        # Load weights
+        weights_path = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin')
+        if os.path.exists(weights_path):
+            state_dict = torch.load(weights_path, map_location='cpu')
+            model.load_state_dict(state_dict)
+        
+        return model
+
+
+# ========== Model Factory ==========
+
+# Mapping từ model key sang Hugging Face model ID
+# Các key này được dùng trong main.py, cần map sang Hugging Face model ID thực tế
+MODEL_ID_MAPPING = {
+    "phobert-v1": "vinai/phobert-base",  # PhoBERT v1 (có thể dùng phobert-base hoặc phobert-base-v1)
+    "phobert-v2": "vinai/phobert-base",  # PhoBERT v2
+    "bartpho": "vinai/bartpho-syllable",
+    "visobert": "uitnlp/visobert",
+    "mbert": "bert-base-multilingual-cased",
+    "vit5": "VietAI/vit5-base",
+    "xlm-r": "xlm-roberta-base",
+    "xlm-roberta": "xlm-roberta-base",
+    "roberta-gru": "roberta-base",
+    "roberta-base-gru": "roberta-base",
+}
+
+def get_hf_model_id(model_name: str) -> str:
+    """
+    Chuyển đổi model key sang Hugging Face model ID
+    
+    Args:
+        model_name: Model key (ví dụ: "phobert-v1") hoặc Hugging Face ID (ví dụ: "vinai/phobert-base")
+    
+    Returns:
+        Hugging Face model ID
+    """
+    model_name_lower = model_name.lower()
+    
+    # Nếu đã là Hugging Face ID (có dấu /), trả về nguyên
+    if '/' in model_name:
+        return model_name
+    
+    # Nếu là key, map sang Hugging Face ID
+    if model_name_lower in MODEL_ID_MAPPING:
+        return MODEL_ID_MAPPING[model_name_lower]
+    
+    # Nếu không có trong mapping, giả định đã là Hugging Face ID
+    return model_name
+
+def get_model_class(model_name: str):
+    """Factory function để lấy model class dựa trên model name"""
+    model_name_lower = model_name.lower()
+    
+    # RoBERTa-GRU (check first to avoid confusion)
+    if 'roberta' in model_name_lower and ('gru' in model_name_lower or 'roberta-base-gru' in model_name_lower):
+        return RoBERTaGRUForABSA
+    
+    # Roberta-based (PhoBERT v1/v2, ViSoBERT)
+    if 'phobert' in model_name_lower or 'visobert' in model_name_lower or 'roberta' in model_name_lower:
+        return TransformerForABSA
+    
+    # XLM-RoBERTa
+    elif 'xlm-roberta' in model_name_lower or 'xlm_roberta' in model_name_lower:
+        return XLMRobertaForABSA
+    
+    # BERT
+    elif 'bert' in model_name_lower and 'roberta' not in model_name_lower:
+        return BERTForABSA
+    
+    # BART
+    elif 'bart' in model_name_lower:
+        return BartForABSA
+    
+    # T5
+    elif 't5' in model_name_lower or 'vit5' in model_name_lower:
+        return T5ForABSA
+    
+    # TextCNN
+    elif 'textcnn' in model_name_lower or 'cnn' in model_name_lower:
+        return TextCNNForABSA
+    
+    # BiLSTM
+    elif 'bilstm' in model_name_lower or 'lstm' in model_name_lower:
+        return BiLSTMForABSA
+    
+    # Default: try Roberta
+    else:
+        return TransformerForABSA
+
+
+def create_model(model_name: str, num_aspects: int, num_sentiments: int, vocab_size=None, **kwargs):
+    """
+    Create model instance dựa trên model name
+    
+    Args:
+        model_name: Tên model hoặc model ID từ Hugging Face
+        num_aspects: Số lượng aspects
+        num_sentiments: Số lượng sentiment classes
+        vocab_size: Vocabulary size (chỉ cần cho TextCNN/BiLSTM)
+        **kwargs: Additional arguments
+    """
+    model_class = get_model_class(model_name)
+    
+    # RoBERTa-GRU cần base model riêng
+    if model_class == RoBERTaGRUForABSA:
+        # Use roberta-base as base model for RoBERTa-GRU
+        base_model_name = 'roberta-base'
+        return model_class.from_pretrained(
+            base_model_name,
+            num_aspects=num_aspects,
+            num_sentiments=num_sentiments,
+            trust_remote_code=True,
+            **kwargs
+        )
+    
+    # Non-transformer models
+    if model_class in [TextCNNForABSA, BiLSTMForABSA]:
+        if vocab_size is None:
+            raise ValueError(f"vocab_size is required for {model_class.__name__}")
+        
+        if model_class == TextCNNForABSA:
+            return TextCNNForABSA(
+                vocab_size=vocab_size,
+                embed_dim=kwargs.get('embed_dim', 300),
+                num_filters=kwargs.get('num_filters', 100),
+                filter_sizes=kwargs.get('filter_sizes', [3, 4, 5]),
+                num_aspects=num_aspects,
+                num_sentiments=num_sentiments,
+                max_length=kwargs.get('max_length', 256)
+            )
+        elif model_class == BiLSTMForABSA:
+            return BiLSTMForABSA(
+                vocab_size=vocab_size,
+                embed_dim=kwargs.get('embed_dim', 300),
+                hidden_dim=kwargs.get('hidden_dim', 256),
+                num_layers=kwargs.get('num_layers', 2),
+                num_aspects=num_aspects,
+                num_sentiments=num_sentiments,
+                dropout=kwargs.get('dropout', 0.3)
+            )
+    
+    # Transformer models
+    else:
+        # Chuyển đổi model key sang Hugging Face model ID
+        hf_model_id = get_hf_model_id(model_name)
+        
+        return model_class.from_pretrained(
+            hf_model_id,
+            num_aspects=num_aspects,
+            num_sentiments=num_sentiments,
+            trust_remote_code=True,
+            **kwargs
+        )