Model Description
A BERT-based multi-label classifier that predicts one or more movie genres from a movie’s title and overview. The model supports 19 genres, including Action, Comedy, Drama, Horror, Romance, Science Fiction, and Thriller.
Usage
Load the model with Hugging Face Transformers and provide a movie title or overview. Apply sigmoid to the model outputs and select genres with a probability of at least 0.5.
from transformers import pipeline
classifier = pipeline(
"text-classification",
model="prashantkhandelwal/movie-genres-classification",
top_k=None,
)
result = classifier(
"A crew travels through deep space to stop an alien threat from destroying Earth."
)
genres = [label for label in result[0] if label["score"] >= 0.5]
print(genres)