Upload model

modeling_vnsabsa.py

@@ -1,6 +1,6 @@
 from transformers import PreTrainedModel
-from modules import SmartphoneBERT
 import torch
+import torch.nn as nn
 
 from .configuration_vnsabsa import VnSmartphoneAbsaConfig
 
@@ -72,4 +72,132 @@ class VnSmartphoneAbsaModel(PreTrainedModel):
             if a_i[-1] >= aspect_thresholds[-1]:
                 res_i["OTHERS"] = ""
         
-        return results
+        return results
+
+
+class AspectClassifier(nn.Module):
+    def __init__(
+        self, 
+        input_size: int,
+        dropout: float = 0.3,
+        hidden_size: int = 64,
+        *args, **kwargs
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.input_size = input_size
+
+        self.fc = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(
+                in_features=input_size,
+                out_features=hidden_size
+            ),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(
+                in_features=hidden_size,
+                out_features=10+1
+            )
+        )
+    
+    def forward(self, input: torch.Tensor):
+        x = self.fc(input)
+        return x
+
+
+class PolarityClassifier(nn.Module):
+    def __init__(
+        self, 
+        input_size: int,
+        dropout: float = 0.5,
+        hidden_size: int = 64,
+        *args, **kwargs
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.polarity_fcs = nn.ModuleList([
+            nn.Sequential(
+                nn.Dropout(dropout),
+                nn.Linear(
+                    in_features=input_size,
+                    out_features=hidden_size
+                ),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(
+                    in_features=hidden_size,
+                    out_features=3
+                )
+            )
+            for _ in torch.arange(10)
+        ])
+
+    def forward(self, input: torch.Tensor):
+        polarities = torch.stack([
+            fc(input)
+            for fc in self.polarity_fcs
+        ])
+        
+        if input.ndim == 2:
+            polarities = polarities.transpose(0, 1)
+        return polarities
+
+
+class SmartphoneBERT(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        embed_dim: int = 768,
+        num_heads: int = 8,
+        num_encoders: int = 4,
+        encoder_dropout: float = 0.1,
+        fc_dropout: float =0.4,
+        fc_hidden_size: int = 128,
+        *args, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.embed = nn.Embedding(
+            num_embeddings=vocab_size,
+            embedding_dim=embed_dim,
+            padding_idx=0
+        )
+        self.encoder = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=embed_dim,
+                nhead=num_heads,
+                dim_feedforward=embed_dim,
+                dropout=encoder_dropout,
+                batch_first=True
+            ),
+            num_layers=num_encoders,
+            norm=nn.LayerNorm(embed_dim),
+            enable_nested_tensor=False
+        )
+        self.a_fc = AspectClassifier(
+            input_size=2*embed_dim,
+            dropout=fc_dropout,
+            hidden_size=fc_hidden_size
+        )
+        self.p_fc = PolarityClassifier(
+            input_size=2*embed_dim,
+            dropout=fc_dropout,
+            hidden_size=fc_hidden_size
+        )
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor
+    ):
+        padding_mask = ~attention_mask.bool()
+        x = self.embed(input_ids)
+        x = self.encoder(x, src_key_padding_mask=padding_mask)
+        x[padding_mask] = 0
+        x = torch.cat([
+            x[..., 0, :],
+            torch.mean(x, dim=-2)
+        ], dim=-1)
+
+        a_logits = self.a_fc(x)
+        p_logits = self.p_fc(x)
+        return a_logits, p_logits