-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict.py
31 lines (25 loc) · 1.05 KB
/
predict.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import torch
def get_predictions(model, data_loader, device):
model = model.eval()
thesis = []
predictions = []
prediction_probs = []
real_values = []
with torch.no_grad():
for d in data_loader:
texts = d['combined_thesis']
input_ids = d['input_ids'].to(device)
attention_mask = d['attention_mask'].to(device)
targets = d['targets'].to(device)
outputs = model(input_ids = input_ids,
attention_mask = attention_mask
)
_, pred = torch.max(outputs, dim = 1)
thesis.extend(texts)
predictions.extend(pred)
prediction_probs.extend(outputs)
real_values.extend(targets)
predictions = torch.stack(predictions).cpu()
prediction_probs = torch.stack(prediction_probs).cpu()
real_values = torch.stack(real_values).cpu()
return thesis, predictions, prediction_probs, real_values