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
Model Details
Model Description
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:
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