Upload model

LegibilityModel.py

@@ -0,0 +1,26 @@
+import torch.nn as nn
+from transformers import VisionEncoderDecoderModel, PreTrainedModel, AutoConfig
+
+class LegibilityModel(PreTrainedModel):
+    def __init__(self):
+        config = AutoConfig.from_pretrained("microsoft/trocr-base-handwritten")
+        super(LegibilityModel, self).__init__(config=config)
+
+        # base model architecture
+        self.model = VisionEncoderDecoderModel(config).encoder
+
+        # change dropout during training
+        self.stack = nn.Sequential(
+            nn.Dropout(0),
+            nn.Linear(768, 768),
+            nn.ReLU(),
+            nn.Dropout(0),
+            nn.Linear(768, 1)
+        )
+
+    def forward(self, img_batch):
+        output = self.model(img_batch)
+        # average the output of the last hidden layer
+        output = output.last_hidden_state.mean(dim=1)
+        scores = self.stack(output)
+        return scores

config.json

@@ -4,6 +4,9 @@
   "architectures": [
     "LegibilityModel"
   ],
+  "auto_map": {
+    "AutoModel": "LegibilityModel.LegibilityModel"
+  },
   "decoder": {
     "_name_or_path": "",
     "activation_dropout": 0.0,