Add constructor functions for model and tokenizer of MonoBERT/T5 (#93)

* Add construct functions for model and tokenizer of MonoBERT/T5

* Remove useless import

* Add pretrained_name as argument

README.md

@@ -38,9 +38,7 @@ Here's how to initalize the T5 reranker from [Document Ranking with a Pretrained
 from pygaggle.rerank.base import Query, Text
 from pygaggle.rerank.transformer import MonoT5
 
-model_name = 'castorini/monot5-base-msmarco'
-tokenizer_name = 't5-base'
-reranker =  MonoT5(model_name, tokenizer_name)
+reranker =  MonoT5()
 ```
 
 Alternatively, here's the BERT reranker from [Passage Re-ranking with BERT](https://arxiv.org/pdf/1901.04085.pdf), which isn't as good as the T5 reranker:
@@ -49,9 +47,7 @@ Alternatively, here's the BERT reranker from [Passage Re-ranking with BERT](http
 from pygaggle.rerank.base import Query, Text
 from pygaggle.rerank.transformer import MonoBERT
 
-model_name = 'castorini/monobert-large-msmarco'
-tokenizer_name = 'bert-large-uncased'
-reranker =  MonoBERT(model_name, tokenizer_name)
+reranker =  MonoBERT()
 ```
 
 Either way, continue with a complete reranking example:

pygaggle/rerank/transformer.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import List, Union
+from typing import List
 
 from transformers import (AutoTokenizer,
                           AutoModelForSequenceClassification,
@@ -26,19 +26,27 @@
 
 class MonoT5(Reranker):
     def __init__(self,
-                 model_name_or_instance: Union[str, T5ForConditionalGeneration] = 'castorini/monot5-base-msmarco',
-                 tokenizer_name_or_instance: Union[str, QueryDocumentBatchTokenizer] = 't5-base'):
-        if isinstance(model_name_or_instance, str):
-            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-            model_name_or_instance = T5ForConditionalGeneration.from_pretrained(model_name_or_instance).to(device).eval()
-        self.model = model_name_or_instance
-
-        if isinstance(tokenizer_name_or_instance, str):
-            tokenizer_name_or_instance = T5BatchTokenizer(AutoTokenizer.from_pretrained(tokenizer_name_or_instance), batch_size=8)
-        self.tokenizer = tokenizer_name_or_instance
-
+                 model: T5ForConditionalGeneration = None,
+                 tokenizer: QueryDocumentBatchTokenizer = None):
+        self.model = model or self.get_model()
+        self.tokenizer = tokenizer or self.get_tokenizer()
         self.device = next(self.model.parameters(), None).device
 
+    @staticmethod
+    def get_model(pretrained_model_name_or_path: str = 'castorini/monot5-base-msmarco',
+                  *args, device: str = None, **kwargs) -> T5ForConditionalGeneration:
+        device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        device = torch.device(device)
+        return T5ForConditionalGeneration.from_pretrained(pretrained_model_name_or_path, *args, **kwargs).to(device).eval()
+
+    @staticmethod
+    def get_tokenizer(pretrained_model_name_or_path: str = 't5-base',
+                      *args, batch_size: int = 8, **kwargs) -> T5BatchTokenizer:
+        return T5BatchTokenizer(
+            AutoTokenizer.from_pretrained(pretrained_model_name_or_path, *args, **kwargs),
+            batch_size=batch_size
+        )
+
     def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
         texts = deepcopy(texts)
         batch_input = QueryDocumentBatch(query=query, documents=texts)
@@ -108,19 +116,24 @@ def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
 
 class MonoBERT(Reranker):
     def __init__(self,
-                 model_name_or_instance: Union[str, PreTrainedModel] = 'castorini/monobert-large-msmarco',
-                 tokenizer_name_or_instance: Union[str, PreTrainedTokenizer] = 'bert-large-uncased'):
-        if isinstance(model_name_or_instance, str):
-            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-            model_name_or_instance = AutoModelForSequenceClassification.from_pretrained(model_name_or_instance).to(device).eval()
-        self.model = model_name_or_instance
-
-        if isinstance(tokenizer_name_or_instance, str):
-            tokenizer_name_or_instance = AutoTokenizer.from_pretrained(tokenizer_name_or_instance)
-        self.tokenizer = tokenizer_name_or_instance
-
+                 model: PreTrainedModel = None,
+                 tokenizer: PreTrainedTokenizer = None):
+        self.model = model or self.get_model()
+        self.tokenizer = tokenizer or self.get_tokenizer()
         self.device = next(self.model.parameters(), None).device
 
+    @staticmethod
+    def get_model(pretrained_model_name_or_path: str = 'castorini/monobert-large-msmarco',
+                  *args, device: str = None, **kwargs) -> AutoModelForSequenceClassification:
+        device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        device = torch.device(device)
+        return AutoModelForSequenceClassification.from_pretrained(pretrained_model_name_or_path, *args, **kwargs).to(device).eval()
+
+    @staticmethod
+    def get_tokenizer(pretrained_model_name_or_path: str = 'bert-large-uncased',
+                      *args, **kwargs) -> AutoTokenizer:
+        return AutoTokenizer.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
     @torch.no_grad()
     def rerank(self, query: Query, texts: List[Text]) -> List[Text]:
         texts = deepcopy(texts)

pygaggle/run/evaluate_document_ranker.py

@@ -80,11 +80,10 @@ def tokenizer_sane(cls, v: str, values, **kwargs):
 
 
 def construct_t5(options: DocumentRankingEvaluationOptions) -> Reranker:
-    device = torch.device(options.device)
-    model = T5ForConditionalGeneration.from_pretrained(options.model,
-                                                       from_tf=options.from_tf).to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(options.model_type)
-    tokenizer = T5BatchTokenizer(tokenizer, options.batch_size)
+    model = MonoT5.get_model(options.model,
+                             from_tf=options.from_tf,
+                             device=options.device)
+    tokenizer = MonoT5.get_tokenizer(options.model_type, batch_size=options.batch_size)
     return MonoT5(model, tokenizer)
 
 
@@ -102,10 +101,8 @@ def construct_transformer(options:
 
 def construct_seq_class_transformer(options: DocumentRankingEvaluationOptions
                                     ) -> Reranker:
-    model = AutoModelForSequenceClassification.from_pretrained(options.model, from_tf=options.from_tf)
-    device = torch.device(options.device)
-    model = model.to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(options.tokenizer_name)
+    model = MonoBERT.get_model(options.model, from_tf=options.from_tf, device=options.device)
+    tokenizer = MonoBERT.get_tokenizer(options.tokenizer_name)
     return MonoBERT(model, tokenizer)
 
 

pygaggle/run/evaluate_kaggle_highlighter.py

@@ -79,9 +79,9 @@ def construct_t5(options: KaggleEvaluationOptions) -> Reranker:
                                  SETTINGS.flush_cache)
     device = torch.device(options.device)
     model = loader.load().to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(
-                    options.model_name, do_lower_case=options.do_lower_case)
-    tokenizer = T5BatchTokenizer(tokenizer, options.batch_size)
+    tokenizer = MonoT5.get_tokenizer(options.model_type,
+                                     do_lower_case=options.do_lower_case,
+                                     batch_size=options.batch_size)
     return MonoT5(model, tokenizer)
 
 
@@ -103,13 +103,13 @@ def construct_transformer(options: KaggleEvaluationOptions) -> Reranker:
 def construct_seq_class_transformer(options:
                                     KaggleEvaluationOptions) -> Reranker:
     try:
-        model = AutoModelForSequenceClassification.from_pretrained(
-                    options.model_name)
+        model = MonoBERT.get_model(options.model_name, device=options.device)
     except OSError:
         try:
-            model = AutoModelForSequenceClassification.from_pretrained(
+            model = MonoBERT.get_model(
                         options.model_name,
-                        from_tf=True)
+                        from_tf=True,
+                        device=options.device)
         except AttributeError:
             # Hotfix for BioBERT MS MARCO. Refactor.
             BertForSequenceClassification.bias = torch.nn.Parameter(
@@ -120,9 +120,9 @@ def construct_seq_class_transformer(options:
                         options.model_name, from_tf=True)
             model.classifier.weight = BertForSequenceClassification.weight
             model.classifier.bias = BertForSequenceClassification.bias
-    device = torch.device(options.device)
-    model = model.to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(
+            device = torch.device(options.device)
+            model = model.to(device).eval()
+    tokenizer = MonoBERT.get_tokenizer(
                     options.tokenizer_name, do_lower_case=options.do_lower_case)
     return MonoBERT(model, tokenizer)
 

pygaggle/run/evaluate_passage_ranker.py

@@ -78,11 +78,10 @@ def tokenizer_sane(cls, v: str, values, **kwargs):
 
 
 def construct_t5(options: PassageRankingEvaluationOptions) -> Reranker:
-    device = torch.device(options.device)
-    model = T5ForConditionalGeneration.from_pretrained(options.model,
-                                                       from_tf=options.from_tf).to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(options.model_type)
-    tokenizer = T5BatchTokenizer(tokenizer, options.batch_size)
+    model = MonoT5.get_model(options.model,
+                             from_tf=options.from_tf,
+                             device=options.device)
+    tokenizer = MonoT5.get_tokenizer(options.model_type, batch_size=options.batch_size)
     return MonoT5(model, tokenizer)
 
 
@@ -101,8 +100,8 @@ def construct_transformer(options:
 def construct_seq_class_transformer(options: PassageRankingEvaluationOptions
                                     ) -> Reranker:
     try:
-        model = AutoModelForSequenceClassification.from_pretrained(
-            options.model, from_tf=options.from_tf)
+        model = MonoBERT.get_model(
+            options.model, from_tf=options.from_tf, device=options.device)
     except AttributeError:
         # Hotfix for BioBERT MS MARCO. Refactor.
         BertForSequenceClassification.bias = torch.nn.Parameter(
@@ -113,9 +112,9 @@ def construct_seq_class_transformer(options: PassageRankingEvaluationOptions
             options.model, from_tf=options.from_tf)
         model.classifier.weight = BertForSequenceClassification.weight
         model.classifier.bias = BertForSequenceClassification.bias
-    device = torch.device(options.device)
-    model = model.to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(options.tokenizer_name)
+        device = torch.device(options.device)
+        model = model.to(device).eval()
+    tokenizer = MonoBERT.get_tokenizer(options.tokenizer_name)
     return MonoBERT(model, tokenizer)