| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| class MovieTrailerModel: | |
| def __init__(self, model_name='bert-base-uncased'): | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| self.model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=5) | |
| def predict(self, prompt): | |
| inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True) | |
| with torch.no_grad(): | |
| outputs = self.model(**inputs) | |
| probabilities = torch.sigmoid(outputs.logits) # Convert logits to probabilities | |
| return probabilities |