Create README.md

README.md

@@ -0,0 +1,14 @@
+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