Upload 2 files

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb98898cc875973d1f028b523405feea9140410d2895dec8dfb675e68c2d08aa
+size 265516836

torch_primitives.py

@@ -0,0 +1,74 @@
+import pandas as pd
+
+import torch
+from torch.utils.data import Dataset
+import torch.nn as nn
+
+from transformers import DistilBertTokenizer, DistilBertModel
+
+
+class PaperClassifierDatasetV1(Dataset):
+    MAJORS = ('cs', 'math', 'physics', 'q-bio', 'q-fin', 'stat', 'econ', 'eess')
+    def __init__(self, csv_path: str, no_abstract_proba: float = 0., n_samples: int = 0):
+        super().__init__()
+        self.major_to_idx = {major : idx for idx, major in enumerate(self.MAJORS)}
+        self.n_classes = len(self.MAJORS)
+
+        self._tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
+
+        self.df = pd.read_csv(csv_path)
+
+        if n_samples == 0:
+            n_samples = self.df.shape[0]
+
+        self.x = self._tokenizer(
+            list(zip(self.df['title'], self.df['abstract']))[:n_samples],
+            padding=True,
+            truncation=True,
+            max_length=512,
+            return_tensors='pt'
+        )
+
+        self.y = torch.zeros((n_samples, len(self.MAJORS)))
+        for row_idx, majors in enumerate(self.df['majors'][:n_samples]):
+            majors = eval(majors)
+            col_idxs = [self.major_to_idx[major] for major in majors]
+            self.y[row_idx, col_idxs] = 1
+
+        self.sep_token_id = self._tokenizer.sep_token_id
+        self.sep_positions = list()
+        for row_idx in range(len(self.x['input_ids'])):
+            input_ids = self.x['input_ids'][row_idx]
+            sep_pos = (input_ids == self.sep_token_id).nonzero(as_tuple=True)[0][0]
+            self.sep_positions.append(sep_pos)
+
+        self.no_abstract_proba = no_abstract_proba
+
+    def __getitem__(self, index: int):
+        input_ids = self.x['input_ids'][index, ...]
+        attention_mask = self.x['attention_mask'][index, ...].clone()
+
+        if torch.rand(1).item() < self.no_abstract_proba:
+            sep_pos = self.sep_positions[index]
+            attention_mask[sep_pos+1:] = 0
+
+        return {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask,
+            'target': self.y[index, ...]
+        }
+
+    def __len__(self):
+        return self.x['input_ids'].shape[0]
+
+
+class PaperClassifierV1(nn.Module):
+    def __init__(self, n_classes: int):
+        super().__init__()
+        self.backbone = DistilBertModel.from_pretrained("distilbert-base-uncased")
+        self.head = nn.Linear(in_features=self.backbone.config.hidden_size, out_features=n_classes)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
+        backbone_output = self.backbone(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
+        logits = self.head(backbone_output[:, 0, ...])
+        return logits