README.md

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