Update models.py

models.py

@@ -1,20 +1,70 @@
-import torch
-import torch.nn as nn
-from transformers import AutoModel, AutoConfig
-
-class TransformerForEmotion(nn.Module):
-    def __init__(self, model_name: str, num_labels: int):
-        super().__init__()
-        config = AutoConfig.from_pretrained(model_name, num_labels=num_labels)
-        self.encoder = AutoModel.from_pretrained(model_name, config=config)
-        self.classifier = nn.Linear(config.hidden_size, num_labels)
-
-    def forward(self, input_ids, attention_mask, labels=None):
-        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
-        pooled = out.last_hidden_state[:, 0]  # CLS token
-        logits = self.classifier(pooled)
-        loss = None
-        if labels is not None:
-            loss_fn = nn.CrossEntropyLoss()
-            loss = loss_fn(logits, labels)
-        return {"loss": loss, "logits": logits}
+import torch
+import os
+import torch.nn as nn
+from transformers import RobertaPreTrainedModel, RobertaModel, AutoConfig
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+class TransformerForABSA(RobertaPreTrainedModel):
+    base_model_prefix = "roberta"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = RobertaModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        # Thêm lớp "none" vào num_sentiments
+        self.sentiment_classifiers = nn.ModuleList([
+            nn.Linear(config.hidden_size, config.num_sentiments + 1)  # +1 cho "none"
+            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)  # [B, H]
+        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)  # [B, A, S+1]
+
+        loss = None
+        if labels is not None:
+            # labels: [B, A], với lớp "none" thay vì -100
+            B, A, _ = all_logits.size()
+            logits_flat = all_logits.view(-1, all_logits.size(-1))
+            targets_flat = labels.view(-1)
+            loss_fct = nn.CrossEntropyLoss()  # Không dùng ignore_index
+            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:]
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=all_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+            """
+            HuggingFace Trainer đôi khi truyền vào state_dict=..., nên ta
+            chấp nhận thêm **kwargs để không vướng lỗi.
+            """
+            # 1) Lưu phần backbone (encoder)
+            self.roberta.save_pretrained(save_directory, **kwargs)  
+
+            # 2) Cập nhật config rồi lưu
+            config = self.roberta.config
+            config.num_aspects = len(self.sentiment_classifiers)
+            config.num_sentiments = self.sentiment_classifiers[0].out_features
+            config.auto_map       = {"AutoModel": "models.TransformerForABSA"}
+            config.save_pretrained(save_directory, **kwargs)
+
+            # 3) Lưu toàn bộ state_dict (bao gồm cả 2 head) — 
+            #    nếu Trainer đã truyền state_dict trong kwargs, có thể dùng luôn,
+            #    nếu không, lấy từ self.state_dict()
+            sd = kwargs.get("state_dict", None) or self.state_dict()
+            torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))