File size: 2,033 Bytes
7b41c88 16d16e2 7b41c88 |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 |
from transformers import RobertaTokenizer, RobertaForSequenceClassification
import torch
import numpy as np
class Model:
def __init__(self, model_weights):
self.tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
self.model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=4)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ✅ Load Lightning checkpoint
checkpoint = torch.load(model_weights, map_location=self.device)
state_dict = checkpoint.get("state_dict", checkpoint)
# ✅ Remove 'model.' prefix used by LightningModule
filtered_state_dict = {
k.replace("model.", ""): v
for k, v in state_dict.items()
if k.startswith("model.")
}
# ✅ Load weights into Hugging Face model
self.model.load_state_dict(filtered_state_dict, strict=False)
self.currepoch = checkpoint.get("epoch", "N/A")
self.loss = checkpoint.get("loss", "N/A")
print(f"✅ Loaded model state — Epoch: {self.currepoch}, Loss: {self.loss}")
self.model.to(self.device)
self.model.eval()
self.labels = ["Blocker", "Critical", "Major", "Minor"]
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
probs = torch.nn.functional.softmax(logits, dim=-1)
predicted_label = self.labels[torch.argmax(probs).item()]
return probs[0].tolist(), predicted_label
# Singleton instance
model_instance = None
model_weights = "assets/roberta-priority-epoch=06-val_f1=0.72.ckpt" # Update path if needed
def get_model():
global model_instance
if model_instance is None:
model_instance = Model(model_weights)
return model_instance
|