Create modeling_roberta.py

modeling_roberta.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+from transformers import (
+    RobertaModel,
+    RobertaForQuestionAnswering,
+)
+
+
+class SpanPredictionHead(nn.Module):
+    """Head for span prediction tasks.
+    Can be viewed as a 2-class output layer that is applied to every position.
+    """
+
+    def __init__(self, input_dim, inner_dim, num_classes, pooler_dropout):
+        assert num_classes == 2
+        super().__init__()
+        self.dense = nn.Linear(input_dim, inner_dim)
+        self.dropout = nn.Dropout(p=pooler_dropout)
+        self.out_proj = nn.Linear(inner_dim, num_classes)
+
+    def forward(self, features, **kwargs):
+        x = features  # take features across ALL positions
+        x = self.dropout(x)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x  # B x T x C, but softmax should be taken over T
+
+
+class RobertaForPororoMRC(RobertaForQuestionAnswering):
+    
+    def __init__(self, config):
+        # Initialize on RobertaPreTrainedModel
+        super(RobertaForQuestionAnswering, self).__init__(config)
+        
+        config.num_labels = 2
+        self.num_labels = config.num_labels
+        
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.qa_outputs = SpanPredictionHead(
+            input_dim=config.hidden_size,
+            inner_dim=config.span_head_inner_dim,
+            num_classes=config.num_labels,
+            pooler_dropout=config.span_head_dropout,
+        )
+        
+        self.init_weights()