This repository contains the code to train a Question-Answering (QA) model based on the ALBERT language model. The training and evaluation is done on the Stanford Question Answering Dataset (SQuAD) version 2.0. The model is designed to answer questions based on a given context and has been fine-tuned for this specific task.
- Environment Setup: The code includes the setup for the Python environment, including the installation of necessary libraries and dependencies using pip.
- Data Preparation: It provides functions for parsing and preparing the SQuAD 2.0 dataset, including reading JSON files and modifying answer contexts for better alignment.
- Custom Dataset: The code defines a custom dataset class for DataLoader, allowing efficient data loading and processing for training and validation.
- Model Configuration: It configures the ALBERT model for Question Answering, setting hyperparameters such as learning rate, batch sizes, and number of epochs.
- Training: The repository includes the training loop for the QA model, along with logging and metrics tracking. It utilizes PyTorch Lightning for streamlined training.
- Learning Rate Finder: There is code to run a learning rate finder, which helps identify an optimal learning rate for training the model.
- Model Checkpointing: It implements model checkpointing, saving the best model based on validation loss.
- Testing: The code includes an example of how to use the trained model to answer questions from a given context.
- Model Saving: It saves the trained model, tokenizer, and training arguments for future use.