Downloads · 30 days
6
4% of all-time downloads
patelnikunjr/distilbert-goodreads-genres
distilbert-goodreads-genres is a text classification model from patelnikunjr. Use it when you need a label for a piece of text. It is set up for transformers. The card lists the license as apache-2.0.
This is a distilbert-base-cased model that is fine-tuned to find the genre of a book from its review text. It can predict 7 genres: comicsgraphic, fantasyparanormal, historybiography, mysterythrillercrime, poetry, rom…
Downloads · 30 days
6
4% of all-time downloads
All-time downloads
145
Public
Parameters
65.8M
1.1 GB on disk
Likes
0
Public
Click a slice to open those files.
.safetensors263 MB · 100%
From the Hugging Face model README
This is a distilbert-base-cased model that is fine-tuned to find the genre of a book from its review text. It can predict 7 genres: comics_graphic, fantasy_paranormal, history_biography, mystery_thriller_crime, poetry, romance, and young_adult.
This model was built as part of MLOps Assignment 2 for the PGD AI Program at IIT Jodhpur. The goal of the assignment is to learn the full MLOps process — loading data, training a model on a GPU, tracking the experiment with Weights & Biases, and uploading the model to Hugging Face.
The model takes a book review (or a short sentence) as input and gives the genre that the book most likely belongs to. It is trained on book reviews from the UCSD Goodreads dataset. The base DistilBERT model is used and a classification layer is added on top of it for this task.
The model can be used directly with the Hugging Face text-classification pipeline. A book review or a short sentence is given as input, and the model returns the predicted genre. It is mainly meant for learning and for the assignment.
This model can be used as a starting point for other projects, like adding genre tags in a reading app or a library system. It can also be trained more on a bigger dataset to make it better.
This model should not be used for text that is not in English. It only knows the 7 genres it was trained on, so it cannot predict any other genre. It should not be used for any important or serious decision.
The training data is taken from public Goodreads reviews, so it can have the biases of the people who wrote those reviews. Also, many books belong to more than one genre at the same time, but this model can only pick one genre. Because of this, the accuracy is not very high (around 58%). Some genres like poetry and comics_graphic work well, but genres like young_adult and fantasy_paranormal get mixed up because they are similar.
Users should treat the result as a suggestion only and not as the final answer. It is better to also check the confidence score. The model should not be used where a wrong genre can cause a problem. The model can be made better by using more data.
Use the code below to get started with the model.
from transformers import pipeline
classifier = pipeline(
"text-classification",
model="patelnikunjr/distilbert-goodreads-genres",
)
result = classifier(
"The detective found a fingerprint on the cold brass handle."
)
print(result)
# Output: [{'label': 'mystery_thriller_crime', 'score': 0.8691}]
The dataset used is the UCSD Goodreads Reviews dataset (from McAuley Lab). It has separate review files for each genre. For each of the 7 genres, the first 10,000 reviews were read and a sample of 2,000 reviews was kept. So in total there were 7,000 reviews. This was split into 5,600 reviews for training and 1,400 reviews for testing (200 test reviews for each genre).
The DistilBertTokenizerFast tokenizer for distilbert-base-cased was used to convert the review text into tokens. The maximum length was set to 512, with truncation and padding. The genre names were converted into numbers using a label map, and then a custom PyTorch Dataset was created.
The Hugging Face Trainer was used for training. Evaluation and a checkpoint were done after every epoch, and load_best_model_at_end=True was set. The final model size is around 263 MB. Testing on 1,400 reviews took around 14 seconds on a single T4 GPU.
The test set has 1,400 reviews (200 reviews per genre). These reviews are different from the training reviews.
The results are checked for each genre separately, because some genres are easier and some are harder to predict.
Accuracy and weighted F1 score are used. Precision, recall and F1 are also checked for each genre. Weighted F1 is used as the main metric because it considers all 7 classes.
Final result: accuracy = 0.585, weighted F1 = 0.585, eval loss = 2.197.
| Genre | Precision | Recall | F1-score | Support |
|---|---|---|---|---|
| comics_graphic | 0.81 | 0.81 | 0.81 | 200 |
| fantasy_paranormal | 0.42 | 0.42 | 0.42 | 200 |
| history_biography | 0.56 | 0.64 | 0.60 | 200 |
| mystery_thriller_crime | 0.50 | 0.57 | 0.53 | 200 |
| poetry | 0.82 | 0.80 | 0.81 | 200 |
| romance | 0.65 | 0.51 | 0.57 | 200 |
| young_adult | 0.37 | 0.35 | 0.36 | 200 |
| Accuracy | 0.585 | 1400 | ||
| Macro avg | 0.59 | 0.58 | 0.59 | 1400 |
| Weighted avg | 0.59 | 0.58 | 0.59 | 1400 |
The model works well for genres that have their own special words, like poetry and comics_graphic (F1 around 0.81). It works less well for genres like young_adult and fantasy_paranormal because these genres are similar and use similar words. The overall accuracy is 0.58, which is much better than random guessing (which would be around 0.14 for 7 classes).
The model was also tested on some new sentences to check if the predictions make sense:
| Input text | Predicted genre | Score |
|---|---|---|
| The detective found a fingerprint on the cold brass handle. | mystery_thriller_crime | 0.8691 |
| Their eyes met across the crowded ballroom and she felt her heart race. | romance | 0.2745 |
| In the quiet of the night, the soul speaks in rhyme. | poetry | 0.9415 |
| The 19th-century industrial revolution reshaped the global economy. | history_biography | 0.4661 |
| The dragon unfurled its wings above the ancient stone castle. | fantasy_paranormal | 0.5355 |
Most of the predictions look correct.
Carbon emissions can be estimated using the Machine Learning Impact calculator presented in Lacoste et al. (2019).
The model is a DistilBERT encoder (6 layers, around 66M parameters) with a classification head on top that gives 7 outputs (one for each genre). The training uses cross-entropy loss.
One NVIDIA Tesla T4 GPU from the Kaggle Notebook environment.
Python 3, PyTorch, Hugging Face Transformers and Datasets, scikit-learn, the huggingface_hub library, and Weights & Biases for experiment tracking.
BibTeX:
@article{sanh2019distilbert,
title={DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter},
author={Sanh, Victor and Debut, Lysandre and Chaumond, Julien and Wolf, Thomas},
journal={arXiv preprint arXiv:1910.01108},
year={2019}
}
APA:
Sanh, V., Debut, L., Chaumond, J., & Wolf, T. (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv:1910.01108.
This model was built for MLOps Assignment 2 in the MTech AI program at IIT Jodhpur. The full code, the data part, and the evaluation part are in the GitHub repository given above. All training runs are saved in the W&B dashboard.
Nikunj R Patel (G25AIT2072)
Nikunj R Patel — through the GitHub repository (G25AIT2072).