---
library_name: transformers
tags:
- gec
- grammar
language:
- en
metrics:
- accuracy
- f1
base_model:
- FacebookAI/roberta-large
pipeline_tag: token-classification
---

# Model Card for Model ID

<!-- Provide a quick summary of what the model is/does. -->

## Model Details

### Model Description

<!-- Provide a longer summary of what this model is. -->

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

<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->

### Direct Use

<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->

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 section addresses misuse, malicious use, and uses that the model will not work well for. -->

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

<!-- This section is meant to convey both technical and sociotechnical 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

<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->

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

<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->

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

<!-- This section describes the evaluation protocols and provides the results. -->

### Testing Data, Factors & Metrics

#### Testing Data

<!-- This should link to a Dataset Card if possible. -->

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

#### Metrics

<!-- These are the evaluation metrics being used, ideally with a description of why. -->

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