Spaces:
Running
Running
File size: 1,638 Bytes
21dc881 3a3df03 21dc881 41fecdc 21dc881 9328f0f 21dc881 3a3df03 21dc881 9716ed4 3a3df03 21dc881 |
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 37 38 39 40 41 42 43 |
from transformers import AlbertTokenizer, AlbertForSequenceClassification
import torch
import numpy as np
class Model:
def __init__(self, model_weights):
self.tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
self.model = AlbertForSequenceClassification.from_pretrained('albert-base-v2', num_labels=5)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the checkpoint
checkpoint = torch.load(model_weights, map_location=self.device)
# Load the model's state dictionary
self.model.load_state_dict(checkpoint['model_state_dict'],strict=False)
self.currepoch = checkpoint['epoch']
self.loss = checkpoint['loss']
print(f"Loaded model state: Current epoch {self.currepoch}, current loss {self.loss}")
self.model.to(self.device)
self.model.eval()
def predict(self, text):
inputs = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = self.model(**inputs)
logits = outputs.logits
predictions = torch.nn.functional.softmax(logits, dim=-1)
Labels = ["No", "Yes"]
return predictions[0].tolist()[:2],Labels[np.argmax(predictions)]
model_instance = None
model_weights = "assets/albert_sentiment_checkpoint_58.pt"
def get_model():
global model_instance
if model_instance is None:
model_instance = Model(model_weights)
return model_instance
|