Add albert result

monologg · Dec 13, 2019 · f4556dd · f4556dd
1 parent 30e9d86
commit f4556dd
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -58,13 +58,16 @@ Run 5 epochs each (No hyperparameter tuning)
 | **ATIS**  | BERT       | 97.87          | 95.46       |
 |           | DistilBERT | 97.54          | 94.89       |
 |           | RoBERTa    | 97.64          | 94.94       |
+|           | ALBERT     | 97.42          | 94.91       |
 | **Snips** | BERT       | 98.29          | 96.05       |
 |           | DistilBERT | 98.42          | 94.10       |
 |           | RoBERTa    | 98.14          | 94.60       |
+|           | ALBERT     | 98.14          | 95.95       |
 
 ## Updates
 
-- 2019/12/03: Add DistilBert and RoBERTa
+- 2019/12/03: Add DistilBert and RoBERTa result
+- 2019/12/14: Add Albert(large v1) result
 
 ## References
 

diff --git a/model.py b/model.py
@@ -1,6 +1,7 @@
 import torch.nn as nn
 from transformers import BertPreTrainedModel, DistilBertConfig, PreTrainedModel, BertModel, DistilBertModel
 from transformers.modeling_distilbert import DistilBertPreTrainedModel
+from transformers.modeling_albert import AlbertPreTrainedModel
 from utils import PRETRAINED_MODEL_MAP
 
 

diff --git a/utils.py b/utils.py
@@ -6,12 +6,13 @@
 import numpy as np
 from sklearn.metrics import f1_score
 from transformers import BertTokenizer, BertConfig, DistilBertConfig, DistilBertTokenizer, RobertaConfig, RobertaTokenizer, \
-    BertModel, DistilBertModel, RobertaModel
+    BertModel, DistilBertModel, RobertaModel, AlbertConfig, AlbertTokenizer, AlbertModel
 
 PRETRAINED_MODEL_MAP = {
     'bert': BertModel,
     'distilbert': DistilBertModel,
-    'roberta': RobertaModel
+    'roberta': RobertaModel,
+    'albert': AlbertModel
 }
 
 from model import JointBERT, JointDistilBERT
@@ -20,12 +21,14 @@
     'bert': (BertConfig, JointBERT, BertTokenizer),
     'distilbert': (DistilBertConfig, JointDistilBERT, DistilBertTokenizer),
     'roberta': (RobertaConfig, JointBERT, RobertaTokenizer),
+    'albert': (AlbertConfig, JointBERT, AlbertTokenizer)
 }
 
 MODEL_PATH_MAP = {
     'bert': 'bert-base-uncased',
     'distilbert': 'distilbert-base-uncased',
-    'roberta': 'roberta-base'
+    'roberta': 'roberta-base',
+    'albert': 'albert-large-v1'
 }