End of training

config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "TimesheetEstimator"
   ],
+  "auto_map": {
+    "AutoConfig": "index.TimesheetEstimatorConfig",
+    "AutoModel": "index.TimesheetEstimator"
+  },
   "encoder_model_name": "distilroberta-base",
   "hidden_size": 768,
   "id2label": {

index.py

@@ -0,0 +1,49 @@
+from transformers import PreTrainedModel, PretrainedConfig, AutoModel
+from transformers.modeling_outputs import ModelOutput
+import torch.nn as nn
+import torch
+
+class TimesheetEstimatorConfig(PretrainedConfig):
+    def __init__(
+        self, 
+        encoder_model_name = "bert-base-cased",
+        hidden_size=768, 
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        self.num_labels = 1
+        self.hidden_size = hidden_size
+        self.encoder_model_name = encoder_model_name
+
+class TimesheetEstimator(PreTrainedModel):
+    config_class = TimesheetEstimatorConfig
+
+    def __init__(self, config: TimesheetEstimatorConfig):
+        super().__init__(config)
+
+        self.encoder = AutoModel.from_pretrained(config.encoder_model_name)
+        self.estimate_layer = nn.Linear(config.hidden_size, config.num_labels)
+        self.loss = nn.MSELoss()
+
+    def forward(self, input_ids, attention_mask, labels=None):
+        encoder_outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        )
+
+        represent_vectors = encoder_outputs[0]
+
+        estimate = self.estimate_layer(represent_vectors[:, 0, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss(estimate, labels.reshape(-1, 1))
+
+        return ModelOutput(
+            loss=loss,
+            logits=estimate,
+        )
+
+
+