Grammar Is Everything: Grammatical Error Correction

from google.colab import userdata

read_access_token = userdata.get('hf_read')
write_access_token = userdata.get('hf_write')

import importlib
import torch, transformers


if '2.3.0' not in torch.__version__:
  !pip install torch==2.3.0
if transformers.__version__!='4.41.2':
  !pip install transformers==4.41.2

if importlib.util.find_spec('datasets') is None:
  !pip install datasets==2.18.0
  !pip install evaluate==0.4.2
  !pip install accelerate -U

!pip install rouge_score

def alter_sentence(example):
  original_sentence = example["part"]
  altered_sentence = "A " + original_sentence
  return {"original_sentence": original_sentence, "altered_sentence": altered_sentence}

import textwrap

def split_paragraph(example):
  return {"part": [part for foo in example["page"] for part in textwrap.wrap(foo, 100)]}

# load the data

from datasets import load_dataset, Dataset, DatasetDict, concatenate_datasets

wiki_dataset = load_dataset('Olimpiada-AI/ro_wiki', token=read_access_token)

# merge splits into a single dataset
wiki_dataset = concatenate_datasets([wiki_dataset["validation"], wiki_dataset["test"]])

# split into smaller chunks
wiki_dataset = wiki_dataset.map(split_paragraph, batched=True, remove_columns="page")

# alter sentences
wiki_dataset = wiki_dataset.map(alter_sentence, batched=False, remove_columns="part")

# split into train and validation
wiki_dataset= wiki_dataset.train_test_split(test_size=0.05)

# load the pre-trained tokenizer and use it to process the data

from transformers import AutoTokenizer
from transformers import DataCollatorWithPadding, DataCollatorForSeq2Seq

tokenizer = AutoTokenizer.from_pretrained("google/mt5-base")


max_length = 256

train_data = wiki_dataset['train'].map(lambda x: {'input_ids': tokenizer(x['altered_sentence'], truncation=True, max_length=max_length)["input_ids"], 'label': tokenizer(x['original_sentence'], truncation=True, max_length=max_length)["input_ids"]}, remove_columns=["original_sentence", "altered_sentence"])
val_data = wiki_dataset['test'].map(lambda x: {'input_ids': tokenizer(x['altered_sentence'], truncation=True, max_length=max_length)["input_ids"], 'label': tokenizer(x['original_sentence'], truncation=True, max_length=max_length)["input_ids"]}, remove_columns=["original_sentence", "altered_sentence"])

data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer)

from transformers import AutoModelForSequenceClassification, TrainingArguments, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments

model = AutoModelForSeq2SeqLM.from_pretrained("google/mt5-base")

# define the evaluation metric

import evaluate
import numpy as np
import sys

bleu = evaluate.load("bleu")
meteor = evaluate.load('meteor')
rouge = evaluate.load('rouge')


def compute_metrics(eval_pred):
    raw_predictions, raw_labels = eval_pred
    predictions = []
    labels = []

    for pred in raw_predictions:
      pred = list(filter(lambda x: x != -100, pred))
      text_predictions = tokenizer.decode(pred, skip_special_tokens=True)
      predictions.append(text_predictions)

    for label in raw_labels:
      label = list(filter(lambda x: x != -100, label))
      text_labels = tokenizer.decode(label, skip_special_tokens=True)
      labels.append([text_labels])

    res_bleu = bleu.compute(predictions=predictions, references=labels)["bleu"]
    res_meteor = meteor.compute(predictions=predictions, references=labels)["meteor"]
    res_rouge = rouge.compute(predictions=predictions, references=labels)["rougeL"]
    return {"bleu": res_bleu, "meteor": res_meteor, "rouge-L": res_rouge}

# define the model and the training configuration

from transformers import AutoModelForSequenceClassification, TrainingArguments, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments

model.generation_config.max_new_tokens = 256

training_args = Seq2SeqTrainingArguments(
    output_dir="baseline_model",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    predict_with_generate=True,
    do_train=True,
    do_eval=True,
    eval_strategy="epoch",
    save_strategy="epoch",
    num_train_epochs=1,
    logging_steps = 1,
    learning_rate=5e-5,
    warmup_steps=5,
    overwrite_output_dir=True,
    save_total_limit=3,
    bf16=True,
    load_best_model_at_end=True,
    push_to_hub=False, #
    hub_strategy="checkpoint",
    hub_token=write_access_token,
    hub_private_repo=True,
    hub_model_id='baseline_model'
)


trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=val_data,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

torch.cuda.empty_cache()
# execute the model training
trainer.train()

# run the trained model on validation split
eval_out = trainer.predict(val_data)
metrics = eval_out.metrics
print(metrics)

# run the trained model on custom data
test_data = [["Acest este o propizite greșita", "Aceasta este o propoziție corectă"],  # "This is an incorrect sentence" -> "This is a correct sentence"
             ["A Ce fdci?", "Ce faci?"],  # "Ce faci?" = "How are you?"
             ["A un test scurt.", "un test scurt."]]  # "un test scurt." = "a short test."

import pandas as pd
df = pd.DataFrame(test_data)
test_data = Dataset.from_pandas(df.rename(columns={0: "input_sentence", 1: "output_sentence"}))
test_data = test_data.map(lambda x: {'input_ids': tokenizer(x['input_sentence'], truncation=True, max_length=max_length)["input_ids"], 'label': tokenizer(x['output_sentence'], truncation=True, max_length=max_length)["input_ids"]}, remove_columns=["input_sentence", "output_sentence"])

eval_out = trainer.predict(test_data)
metrics = eval_out.metrics
print(metrics)