kolodkin/imdb-genre-distilbert

IMDb Plot → Genre Classifier

Fine-tuned distilbert-base-uncased for multi-label genre classification on movie/show plot summaries from kolodkin/imdb-wikipedia-enriched.

Labels

Drama, Comedy, Documentary, Action, Romance, Thriller, Crime, Horror, Adventure, Mystery, Family, Biography, Fantasy, History, Sci-Fi, Music, Animation, War, Musical, Sport, Western

Test results

metric	value
macro-F1	0.4584
micro-F1	0.6070
samples-F1	0.6216

95% bootstrap CI on macro-F1: [0.4507, 0.4658] (1000 resamples).

Per-genre (test)

genre	precision	recall	f1	support
Drama	0.7349	0.7718	0.7529	9889
Comedy	0.7374	0.5142	0.6059	5193
Documentary	0.8816	0.8406	0.8606	3312
Action	0.6874	0.4888	0.5714	2326
Romance	0.5998	0.2889	0.3900	2288
Thriller	0.5400	0.2360	0.3285	2059
Crime	0.5803	0.4018	0.4748	2006
Horror	0.7423	0.6681	0.7032	1901
Adventure	0.6003	0.2956	0.3962	1326
Mystery	0.4583	0.1043	0.1700	949
Family	0.5820	0.2483	0.3480	858
Biography	0.5354	0.2458	0.3369	769
Fantasy	0.5030	0.1125	0.1838	738
History	0.5376	0.1455	0.2291	639
Sci-Fi	0.6510	0.4435	0.5276	593
Music	0.6416	0.4563	0.5333	561
Animation	0.7136	0.3333	0.4544	426
War	0.5925	0.4251	0.4950	407
Musical	0.6452	0.0563	0.1036	355
Sport	0.5799	0.4262	0.4913	298
Western	0.7409	0.6111	0.6698	234

Usage

from transformers import pipeline

clf = pipeline("text-classification", model="kolodkin/imdb-genre-distilbert", top_k=None)
clf("A young wizard discovers a magical school of witchcraft.")

Training

• Multi-label stratified 80 / 10 / 10 split (seed 42), deduped on normalized plot hash
• 5 epochs, lr 2e-05, batch size 128, max length 256
• Best checkpoint selected by val macro-F1
• Source notebook: https://github.com/kolodkin/samples/blob/main/imdb-genre-distilbert/notebook.ipynb