Liwen Huang
12/07/2022, 5:49 PMclass HackathonModel(object):
def __init__(self):
"""Will overwrite model loading here"""
self.tokenizer = RobertaTokenizer.from_pretrained('SkolkovoInstitute/roberta_toxicity_classifier')
self.model = RobertaForSequenceClassification.from_pretrained('SkolkovoInstitute/roberta_toxicity_classifier')
def predict(self, text) -> float:
"""Will overwrite preprocessing and prediction logic here"""
pt_tokens = self.tokenizer.encode(text, return_tensors='pt')
pt_batch = self.tokenizer.encode(inputs[-1], return_tensors='pt')
output_tensor = self.model(pt_batch)
prob_tensor = nn.functional.softmax(output_tensor.logits, dim=-1)[:, 1]
return prob_tensor.detach().numpy()[0]