File size: 1,018 Bytes
a87c588 |
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 32 33 34 35 36 |
import torch
import torch.nn.functional as F
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class ProductionModel():
def __init__(self, tokenizer, dict_labels):
self.model = None
self.tokenizer = tokenizer
self.dict_labels = dict_labels
def predict(self, X):
'''
Método que genera la predicción sobre nuevos datos (X).
X: Lista con los datos, cada elemento es una observación (list)
'''
if self.model is None:
raise ValueError('Debes cargar el modelo con self.model = torch.load(model_file.pt)')
X = self.tokenizer.tokenize(X)
X = torch.tensor(X, device = device)
self.model.eval()
with torch.no_grad():
predictions = self.model(X)[0]
predictions = F.softmax(predictions, dim = 1)
predictions = predictions.to('cpu').detach().numpy()
output = [{self.dict_labels[i]: float(lista[i]) for i in range(len(lista))} for lista in predictions]
if len(output) == 1:
return output[0]
return output |