def predict_output(
self,
sentence: str,
*addl_sentences,
add_special_tokens: bool = True,
no_separator: bool = False,
show_probs: bool = False,
) -> Union[str, Dict]:
tokens = self.encode(
sentence,
*addl_sentences,
add_special_tokens=add_special_tokens,
no_separator=no_separator,
)
with torch.no_grad():
prediction = self.predict(
"sentence_classification_head",
tokens,
return_logits=self.args.regression_target,
)
if self.args.regression_target:
return prediction.item() # float
label_fn = lambda label: self.task.label_dictionary.string(
[label + self.task.label_dictionary.nspecial])
if show_probs:
probs = softmax(prediction.cpu().numpy())
probs = probs.tolist()
probs = {label_fn(i): prob for i, prob in enumerate(probs)}
return probs
return label_fn(prediction.argmax().item()) # str
def predict_output(
self,
sentence: str,
*addl_sentences,
add_special_tokens: bool = True,
no_separator: bool = False,
show_probs: bool = False,
) -> Union[str, Dict]:
"""Predict output, either a classification label or regression target,
using a fine-tuned sentence prediction model.
:returns output
str (classification) or float (regression)
>>> from brain_bert import BrainRobertaModel
>>> model = BrainRobertaModel.load_model('brainbert.base.ko.kornli')
>>> model.predict_output(
... 'BrainBert는 한국어 코퍼스에 학습된 언어모델이다.',
... 'BrainBert는 한국어 모델이다.',
... )
entailment
>>> model = BrainRobertaModel.load_model('brainbert.base.ko.korsts')
>>> model.predict_output(
... 'BrainBert는 한국어 코퍼스에 학습된 언어모델이다.',
... 'BrainBert는 한국어 모델이다.',
... )
0.8374465107917786
"""
assert self.args.task == "sentence_prediction", (
"predict_output() only works for sentence prediction tasks.\n"
"Use predict() to obtain model outputs; "
"use predict_span() for span prediction tasks.")
assert (
"sentence_classification_head" in self.model.classification_heads
), "need pre-trained sentence_classification_head to make predictions"
tokens = self.encode(
sentence,
*addl_sentences,
add_special_tokens=add_special_tokens,
no_separator=no_separator,
)
with torch.no_grad():
prediction = self.predict(
"sentence_classification_head",
tokens,
return_logits=self.args.regression_target,
)
if self.args.regression_target:
return prediction.item() # float
label_fn = lambda label: self.task.label_dictionary.string(
[label + self.task.label_dictionary.nspecial])
if show_probs:
probs = softmax(prediction.cpu().numpy())
probs = probs.tolist()
probs = {label_fn(i): prob for i, prob in enumerate(probs)}
return probs
return label_fn(prediction.argmax().item()) # str
def predict_output(
self,
sentence: str,
*addl_sentences,
no_separator: bool = False,
show_probs: bool = False,
) -> Union[str, float]:
assert self.args.task == "sentence_prediction", (
"predict_output() only works for sentence prediction tasks.\n"
"Use predict() to obtain model outputs; "
"use predict_span() for span prediction tasks.")
assert (
"sentence_classification_head" in self.model.classification_heads
), "need pre-trained sentence_classification_head to make predictions"
tokens = self.encode(
sentence,
*addl_sentences,
no_separator=no_separator,
)
with torch.no_grad():
prediction = self.predict(
"sentence_classification_head",
tokens,
return_logits=self.args.regression_target,
)
if self.args.regression_target:
return prediction.item() # float
label_fn = lambda label: self.task.label_dictionary.string(
[label + self.task.label_dictionary.nspecial])
if show_probs:
probs = softmax(prediction.cpu().numpy())
probs = probs.tolist()
probs = {label_fn(i): prob for i, prob in enumerate(probs)}
return probs
return label_fn(prediction.argmax().item()) # str
