---
license: apache-2.0
base_model: vennify/t5-base-grammar-correction
tags:
- grammar-correction
- t5
- explainable-ai
- text2text-generation
datasets:
- jfleg
widget:
- text: "gec: Hey how are you"
  example_title: "Basic Correction"
---

# Grammar Error Correction (GEC) with T5

This notebook demonstrates how to fine-tune and use a T5-based model for Grammar Error Correction (GEC). It leverages the `transformers` library from Hugging Face for model handling and training.

## Table of Contents

- [Introduction](#introduction)
- [Dataset](#dataset)
- [Model Training](#model-training)
- [Explainable AI Judge](#explainable-ai-judge)
- [Evaluation](#evaluation)
- [Save and Download Model](#save-and-download-model)

## Introduction

This project focuses on building a Grammar Error Correction system using a pre-trained T5 model. GEC is the task of identifying and correcting grammatical errors in text. The notebook covers:

1.  **Data Preparation**: Extracting and parsing M2 format datasets.
2.  **Model Loading**: Loading a T5 model and tokenizer from Hugging Face.
3.  **Fine-tuning**: Training the T5 model on the prepared GEC dataset.
4.  **Explainable AI**: Implementing a custom `ExplainableAIJudge` to provide rationales for corrections.
5.  **Evaluation**: Setting up a basic evaluation framework for GEC and explanation quality.
6.  **Model Export**: Saving and downloading the fine-tuned model.

## Dataset

The notebook uses data from two sources, primarily targeting M2 format files:

-   **CoNLL-14 Shared Task Data**: Used for training (`conll14st-test-data.tar.gz`).
-   **WI+LOCNESS M2 Data (BEA-19)**: Also used for training (`wi+locness_v2.1.bea19.tar.gz`).

The `parse_m2` function extracts source sentences and their corresponding target corrections from these files. The data is then transformed into a format suitable for sequence-to-sequence models.

## Model Training

The core of the GEC system is a T5 model. Specifically, it uses `vennify/t5-base-grammar-correction` as the base model.

### `GECDataset` Class

A custom `GECDataset` class prepares the data for the T5 model, tokenizing source and target sentences and ensuring they are padded/truncated to a maximum length. Each source sentence is prefixed with `"gec: "` to prompt the T5 model for grammar correction.

### `Trainer` from Hugging Face

The `Trainer` API from `transformers` is used for fine-tuning the T5 model. It handles the training loop, logging, and model saving.

## Explainable AI Judge

An `ExplainableAIJudge` class is implemented to not only correct grammar but also provide human-readable explanations for the changes. It leverages `difflib.SequenceMatcher` to find differences between the original and corrected sentences and maps these differences to predefined error types with explanations.

### Error Mapping

The `error_map` dictionary translates internal error codes (e.g., `R:VERB:TENSE`, `R:PUNCT`) into descriptive explanations.

## Evaluation

The `evaluate_xgec` function provides metrics for both correction performance and explanation quality:

-   **Correction Metrics**: Accuracy, Precision, Recall, and F0.5 score, comparing hypotheses (model's corrections) against references (ground truth).
-   **Explanation Quality**: BERTScore F1 for semantic similarity between generated explanations and reference explanations, and Error Type Accuracy for how well the model's identified error types match the ground truth.

## Save and Download Model

After training, the fine-tuned model and tokenizer are saved to a local directory. This directory is then compressed into a `gec_model.tar.gz` file, which can be downloaded using `google.colab.files.download` for deployment or further use.