add models

app.py

@@ -3,22 +3,109 @@ import gradio as gr
 
 if gr.NO_RELOAD:
     import numpy as np
-    from src.distilbert_tf import DistilBertTransferLearningModel
+    from src.model import BaseTransferLearningModel
 
 
 DEVICE = 'cpu'
 MODELS = [
     (
-        'distilbert-1linear-1650',
-        lambda: DistilBertTransferLearningModel(
+        'bert-model_1950',
+        lambda: BaseTransferLearningModel(
+            'bert-base-uncased',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/bert-model_1950.pt',
+        ),
+    ),
+    (
+        'bert-model_2000',
+        lambda: BaseTransferLearningModel(
+            'bert-base-uncased',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/bert-model_2000.pt',
+        ),
+    ),
+    (
+        'deberta-base-model_4400',
+        lambda: BaseTransferLearningModel(
+            'microsoft/deberta-base',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/deberta-base-model_4400.pt',
+        ),
+    ),
+    (
+        'deberta-base-model_8000',
+        lambda: BaseTransferLearningModel(
+            'microsoft/deberta-base',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/deberta-base-model_8000.pt',
+        ),
+    ),
+    (
+        'deberta-v3-base-model_3400',
+        lambda: BaseTransferLearningModel(
+            'microsoft/deberta-v3-base',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/deberta-v3-base-model_3400.pt',
+        ),
+    ),
+    (
+        'deberta-v3-base-model_4000',
+        lambda: BaseTransferLearningModel(
+            'microsoft/deberta-v3-base',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/deberta-v3-base-model_4000.pt',
+        ),
+    ),
+    (
+        'distilbert-model_1850',
+        lambda: BaseTransferLearningModel(
             'distilbert-base-uncased',
-            [
-                ('linear', ['in', 'out']),
-                ('softmax'),
-            ],
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/distilbert-model_1850.pt',
+        ),
+    ),
+    (
+        'distilbert-model_2000',
+        lambda: BaseTransferLearningModel(
+            'distilbert-base-uncased',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/distilbert-model_2000.pt',
+        ),
+    ),
+    (
+        'roberta-base-model_1250',
+        lambda: BaseTransferLearningModel(
+            'FacebookAI/roberta-base',
+            [('linear', ['in', 'out']), ('softmax')],
+            2,
+            device=DEVICE,
+            state_dict='src/ckpt/roberta-base-model_1250.pt',
+        ),
+    ),
+    (
+        'roberta-base-model_2000',
+        lambda: BaseTransferLearningModel(
+            'FacebookAI/roberta-base',
+            [('linear', ['in', 'out']), ('softmax')],
             2,
             device=DEVICE,
-            state_dict='src/ckpt/distilbert-1linear-dataset-all-augmented-all-1650.pt',
+            state_dict='src/ckpt/roberta-base-model_2000.pt',
         ),
     ),
 ]

src/model.py

@@ -0,0 +1,72 @@
+from typing import Any, Optional, Tuple, Union
+import torch
+import transformers
+
+
+class BaseTransferLearningModel(torch.nn.Module):
+
+    def __init__(
+        self,
+        pretrained_model: str = "distilbert-base-uncased",
+        layers: list[Tuple[str, Optional[list[Any]]]] = [
+            ('linear', ['in', 'out']),
+            ('softmax'),
+        ],
+        dim_out: int = 2,
+        use_local_file: bool = False,
+        device: str = 'cpu',
+        state_dict: Optional[Union[str, dict]] = None,
+    ):
+        super(BaseTransferLearningModel, self).__init__()
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(
+            pretrained_model, local_files_only=use_local_file
+        )
+        self.base_model = transformers.AutoModel.from_pretrained(
+            pretrained_model, local_files_only=use_local_file
+        )
+        clf_layers = []
+        for layer in layers:
+            layer_type = layer[0] if isinstance(layer, tuple) else layer
+            if layer_type == 'linear':
+                layer_in, layer_out = [
+                    (
+                        self.base_model.config.hidden_size
+                        if x == 'in'
+                        else dim_out if x == 'out' else x
+                    )
+                    for x in layer[1]
+                ]
+                clf_layers.append(torch.nn.Linear(layer_in, layer_out))
+            elif layer_type == 'softmax':
+                clf_layers.append(torch.nn.Softmax(dim=-1))
+        self.clf = torch.nn.Sequential(*clf_layers)
+
+        if state_dict is not None:
+            if isinstance(state_dict, str) and state_dict.endswith('.pt'):
+                if device == 'cpu':
+                    state_dict = torch.load(state_dict, map_location='cpu')
+                else:
+                    state_dict = torch.load(state_dict)
+            self.load_state_dict(state_dict)
+
+    def forward(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        y = self.base_model(ids, attention_mask=mask, return_dict=False)[0][:, 0]
+        y = self.clf(y)
+        return y
+
+    def predict(self, text: str, device: str) -> torch.Tensor:
+        encoded = self.tokenizer.encode_plus(
+            text,
+            add_special_tokens=True,
+            return_token_type_ids=False,
+            return_attention_mask=True,
+            max_length=512,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt',
+        )
+        with torch.no_grad():
+            ids = encoded['input_ids'].to(device)
+            mask = encoded['attention_mask'].to(device)
+            output = self.forward(ids, mask)
+        return output.to(device)