Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 4 commits into from
Oct 23, 2020
Merged

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

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
16cf45e
Add construct functions for model and tokenizer of MonoBERT/T5
yuxuan-ji Sep 17, 2020
0c0bf88
Remove useless import
yuxuan-ji Sep 17, 2020
9a99f6a
Add pretrained_name as argument
yuxuan-ji Sep 17, 2020
5d0d40d
Merge branch 'master' into simplify-boilerplate-2
yuxuan-ji Oct 23, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 2 additions & 6 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
59 changes: 36 additions & 23 deletions pygaggle/rerank/transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from copy import deepcopy
from typing import List, Union
from typing import List

from transformers import (AutoTokenizer,
AutoModelForSequenceClassification,
Expand Down Expand Up @@ -27,19 +27,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()
Comment on lines +31 to +32
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure but is such code standard python practice?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ronakice, is your concern about the use of the or operand? I've seen this being used as a shorthand to

self.model = model if model else self.get_model()

so I don't think there's any issues there. Or is it more about the PR in general?

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)
Expand Down Expand Up @@ -109,19 +117,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)
Expand Down
15 changes: 6 additions & 9 deletions pygaggle/run/evaluate_document_ranker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand All @@ -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)


Expand Down
20 changes: 10 additions & 10 deletions pygaggle/run/evaluate_kaggle_highlighter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand All @@ -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(
Expand All @@ -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)

Expand Down
19 changes: 9 additions & 10 deletions pygaggle/run/evaluate_passage_ranker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand All @@ -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(
Expand All @@ -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)


Expand Down