This model is a grammar error correction (GEC) system fine-tuned from the FacebookAI/roberta-large model, designed to detect and correct grammatical errors in English text. The model focuses on common grammatical mistakes such as verb tense, noun inflection, adjective usage, and more. It is particularly useful for language learners or applications requiring enhanced grammatical precision.

Model type: Token classification with sequence-to-sequence correction
Language(s) (NLP): English
Finetuned from model: FacebookAI/roberta-large

Uses

Direct Use

This model can be used directly for grammar error detection and correction in English texts. It's ideal for integration into writing assistants, educational software, or proofreading tools.

Downstream Use

The model can be fine-tuned for specific domains like academic writing, business communication, or informal text correction, ensuring high precision in context-specific grammar errors.

Out-of-Scope Use

This model is not suitable for non-English text, non-grammatical corrections (e.g., style, tone, or logic), or detecting complex errors beyond basic grammar structures.

Bias, Risks, and Limitations

The model is trained on general English corpora and may underperform with non-standard dialects (e.g Spoken language), or domain-specific jargon. Users should be cautious when applying it to such contexts, as it might introduce or overlook errors due to the limitations in its training data.

Recommendations

While the model provides strong general performance, users should manually review corrections, especially in specialized or creative contexts where grammar rules can be more fluid.

How to Get Started with the Model

Use the code below to get started with the model.

Use the following code to get started with the model:

python

from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import AutoConfig, AutoTokenizer
from transformers.file_utils import ModelOutput
from transformers.models.roberta.modeling_roberta import (
    RobertaModel,
    RobertaPreTrainedModel,
)

@dataclass
class XGECToROutput(ModelOutput):
    """
    Output type of `XGECToRForTokenClassification.forward()`.
    loss (`torch.FloatTensor`, optional)
    logits_correction (`torch.FloatTensor`) : The correction logits for each token.
    logits_detection (`torch.FloatTensor`) : The detection logits for each token.
    hidden_states (`Tuple[torch.FloatTensor]`, optional)
    attentions (`Tuple[torch.FloatTensor]`, optional)
    """

    loss: Optional[torch.FloatTensor] = None
    logits_correction: torch.FloatTensor = None
    logits_detection: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


class XGECToRRoberta(RobertaPreTrainedModel):
    """
    This class overrides the GECToR model to include an error detection head in addition to the token classification head.
    """

    _keys_to_ignore_on_load_unexpected = [r"pooler"]
    _keys_to_ignore_on_load_missing = [r"position_ids"]

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.unk_tag_idx = config.label2id.get("@@UNKNOWN@@", None)

        self.roberta = RobertaModel(config)

        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        if self.unk_tag_idx is not None:
            self.error_detector = nn.Linear(config.hidden_size, 3)
        else:
            self.error_detector = nn.Linear(config.hidden_size, 2)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        outputs = self.roberta(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]

        logits_correction = self.classifier(sequence_output)
        logits_detection = self.error_detector(sequence_output)

        loss = None
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(
                logits_correction.view(-1, self.num_labels), labels.view(-1)
            )

            labels_detection = torch.ones_like(labels)
            labels_detection[labels == 0] = 0
            labels_detection[labels == -100] = -100  # ignore padding
            if self.unk_tag_idx is not None:
                labels_detection[labels == self.unk_tag_idx] = 2
                loss_detection = loss_fct(
                    logits_detection.view(-1, 3), labels_detection.view(-1)
                )
            else:
                loss_detection = loss_fct(
                    logits_detection.view(-1, 2), labels_detection.view(-1)
                )

            loss += loss_detection

        if not return_dict:
            output = (
                logits_correction,
                logits_detection,
            ) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return XGECToROutput(
            loss=loss,
            logits_correction=logits_correction,
            logits_detection=logits_detection,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def get_input_embeddings(self):
        return self.roberta.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.roberta.set_input_embeddings(value)

config = AutoConfig.from_pretrained("manred1997/roberta-large_lemon-spell_5k")
tokenizer = AutoTokenizer.from_pretrained("manred1997/roberta-large_lemon-spell_5k")
model = XGECToRRoberta.from_pretrained(
    "manred1997/roberta-large_lemon-spell_5k", config=config
)

Training Details

Training Data

We trained the model in three stages, each requiring specific datasets. Below is a description of the data used in each stage:

Stage	Dataset(s) Used	Description
Stage 1	Shuffled 9 million sentences from the PIE corpus (A1 part only)	9 million shuffled sentences from the PIE corpus, focusing on A1-level sentences.
Stage 2	Shuffled combination of NUCLE, FCE, Lang8, W&I + Locness datasets	Lang8 dataset contained 947,344 sentences, with 52.5% having different source and target sentences.
		If using a newer Lang8 dump, consider sampling.
Stage 3	Shuffled version of W&I + Locness datasets	Final shuffled version of the W&I + Locness datasets.

Evaluation

Testing Data, Factors & Metrics

Testing Data

The model was tested on the W&I + Locness test set, a standard benchmark for grammar error correction.

Metrics

The primary evaluation metric used was F0.5, measuring the model's ability to identify and fix grammatical errors correctly.

Results

F0.5 = 73.79