patelnikunjr/distilbert-goodreads-genres
Model Card for distilbert-goodreads-genres
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.
Model Details
Model Description
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.
- Developed by: Nikunj R Patel (G25AIT2072)
- Funded by : IIT Jodhpur
- Shared by : Nikunj R Patel
- Model type: Text classification model (transformer based)
- Language(s) (NLP): English (en)
- License: Apache 2.0
- Finetuned from model : distilbert-base-cased
Model Sources
- Repository: https://github.com/G25AIT2072/MLOPs-Assignment
- Hugging Face Model: https://huggingface.co/patelnikunjr/distilbert-goodreads-genres
- W&B Dashboard: https://wandb.ai/g25ait2072-nikunj/mlops-assignment2
- Paper : DistilBERT paper by Sanh et al. (2019) — https://arxiv.org/abs/1910.01108
- Demo : The code given below can be used to try the model.
Uses
Direct Use
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.
Downstream Use
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.
Out-of-Scope Use
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.
Bias, Risks, and Limitations
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.
Recommendations
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.
How to Get Started with the Model
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}]Training Details
Training Data
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).
Training Procedure
Preprocessing
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.
Training Hyperparameters
- Training regime: fp32
- Epochs: 3
- Train batch size: 16
- Eval batch size: 32
- Learning rate: 3e-5
- Warmup steps: 100
- Weight decay: 0.01
- Max sequence length: 512
- Optimizer: AdamW (default from Hugging Face Trainer)
Speeds, Sizes, Times
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.
Evaluation
Testing Data, Factors & Metrics
Testing Data
The test set has 1,400 reviews (200 reviews per genre). These reviews are different from the training reviews.
Factors
The results are checked for each genre separately, because some genres are easier and some are harder to predict.
Metrics
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.
Results
Final result: accuracy = 0.585, weighted F1 = 0.585, eval loss = 2.197.
Summary
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).
Model Examination
The model was also tested on some new sentences to check if the predictions make sense:
Most of the predictions look correct.
Environmental Impact
Carbon emissions can be estimated using the Machine Learning Impact calculator presented in Lacoste et al. (2019).
- Hardware Type: NVIDIA Tesla T4 GPU (16 GB)
- Hours used: Around 0.3 to 0.5 hours
- Cloud Provider: Kaggle
- Compute Region: Kaggle managed (cannot be selected by user)
- Carbon Emitted: Very low, because it was a short single GPU run
Technical Specifications
Model Architecture and Objective
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.
Compute Infrastructure
Hardware
One NVIDIA Tesla T4 GPU from the Kaggle Notebook environment.
Software
Python 3, PyTorch, Hugging Face Transformers and Datasets, scikit-learn, the huggingface_hub library, and Weights & Biases for experiment tracking.
Citation
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.
Glossary
- Fine-tuning — training a model that is already trained, on a new dataset for a new task.
- Weighted F1 — the F1 score for all classes, given as an average based on the number of samples in each class.
- Token — a small part of the text made by the tokenizer. Here the maximum is 512 tokens.
More Information
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.
Model Card Authors
Nikunj R Patel (G25AIT2072)
Model Card Contact
Nikunj R Patel — through the GitHub repository (G25AIT2072).
