Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3tokenizer = AutoTokenizer.from_pretrained("AdamCodd/distilroberta-query-wellformedness")
4
5class RegressionModel(torch.nn.Module):
6 def __init__(self):
7 super().__init__()
8 self.model = AutoModelForSequenceClassification.from_pretrained("AdamCodd/distilroberta-query-wellformedness")
9 self.regression_head = torch.nn.Linear(self.model.config.hidden_size, 1)
10
11 def forward(self, input_ids, attention_mask, **kwargs):
12 outputs = self.model.base_model(input_ids=input_ids, attention_mask=attention_mask)
13 rating = self.regression_head(outputs.last_hidden_state[:, 0, :])
14 rating = torch.sigmoid(rating)
15 return rating.squeeze()
16
17regression_model = RegressionModel()
18# Do not forget to set the correct path to load the regression head
19regression_model.regression_head.load_state_dict(torch.load("path_to_the_regression_head.pth"))
20regression_model.eval()
21# Examples
22sentences = [
23 "The cat and dog in the yard.",
24 "she don't like apples.",
25 "Is rain sunny days sometimes?",
26 "She enjoys reading books and playing chess.",
27 "How many planets are there in our solar system?"
28]
29
30inputs = tokenizer(sentences, truncation=True, padding=True, return_tensors='pt')
31
32with torch.no_grad():
33 outputs = regression_model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
34
35predictions = outputs.tolist()
36for i, rating in enumerate(predictions):
37 print(f'Sentence: {sentences[i]}')
38 print(f'Predicted Rating: {rating}\n')Sentence: The cat and dog in the yard.
Predicted Rating: 0.20430190861225128
Sentence: she don't like apples.
Predicted Rating: 0.08289700001478195
Sentence: Is rain sunny days sometimes?
Predicted Rating: 0.20011138916015625
Sentence: She enjoys reading books and playing chess.
Predicted Rating: 0.8915354013442993
Sentence: How many planets are there in our solar system?
Predicted Rating: 0.974799394607544'test_loss': 0.061837393790483475,
'test_mse': 0.061837393790483475,
'test_r2': 0.5726782083511353,
'test_mae': 0.183049738407135