synth_translation.py

import transformers
import numpy as np
from datasets import load_dataset, DatasetDict
from transformers import AutoModelForSeq2SeqLM
from transformers import AutoTokenizer
from transformers import DataCollatorForSeq2Seq
import evaluate
import numpy as np
from transformers import Seq2SeqTrainingArguments
from transformers import Seq2SeqTrainer
from torch.utils.data import DataLoader
from transformers import pipeline
from transformers import AdamW
from accelerate import Accelerator
from transformers import get_scheduler
from huggingface_hub import Repository, get_full_repo_name
from tqdm.auto import tqdm
import torch
from torch import Tensor
import os

#load in dataset, setup tokenizer

def addperiod(entry):
    entry['en'] += '.'
    entry['fr'] += '.'
    return entry

raw_datasets = load_dataset("aatherton2024/eng-nah-svo")
train_ds = raw_datasets['train'].map(addperiod)
validation_ds = raw_datasets['validation'].map(addperiod)
test_ds = raw_datasets['test'].map(addperiod)

raw_datasets = DatasetDict({"train" : train_ds, "validation" : validation_ds, "test" : test_ds})
model_checkpoint = "eng-nah-svo-cpt"

if False: #data processing only needs to run once
    def get_training_corpus(raw_datasets):
        return (
            raw_datasets["train"][i : i + 1000]
            for i in range(0, len(raw_datasets["train"]), 1000)
        )

    training_corpus = get_training_corpus(raw_datasets)
    old_tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer = old_tokenizer.train_new_from_iterator(training_corpus, 52000)

    tokenizer.save_pretrained("eng-nah-svo-cpt")
    tokenizer.push_to_hub("eng-nah-svo-cpt")

max_length = 128
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
tokenizer.add_special_tokens({'pad_token': '[PAD]'})

#scan dataset, storing lists of english and french words then returning the tokenization of them
def preprocess_function(examples):
    inputs = examples["en"]
    targets = examples["fr"]
    model_inputs = tokenizer(
        inputs, text_target=targets, max_length=max_length, truncation=True
    )
    return model_inputs

#apply preprocessing in one go to all splits of the dataset
tokenized_datasets = raw_datasets.map(
    preprocess_function,
    batched=True,
    remove_columns=raw_datasets["train"].column_names
)

# #model choice for this problem
if False: #load pretrained model
    model = AutoModelForSeq2SeqLM.from_pretrained("eng-nah-svo-translation")

else:
    from transformers import BertConfig, BertLMHeadModel
    from transformers import AutoModel

    model = AutoModelForSeq2SeqLM.from_pretrained("eng-nah-svo-translation")



#data collator takes tokenizer and the model to deal with padding for dynamic batching
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

#Using BLEU as our metric for this problem
metric_bleu = evaluate.load("sacrebleu")
metric_chrf = evaluate.load("chrf")

#simple method to return test metrics
def compute_metrics(eval_preds):
    preds, labels = eval_preds
    # In case the model returns more than the prediction logits
    if isinstance(preds, tuple):
        preds = preds[0]

    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

    # Replace -100s in the labels as we can't decode them
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    # Some simple post-processing
    decoded_preds = [pred.strip() for pred in decoded_preds]
    decoded_labels = [[label.strip()] for label in decoded_labels]

    result_bleu = metric_bleu.compute(predictions=decoded_preds, references=decoded_labels)
    result_chrf = metric_chrf.compute(predictions=decoded_preds, references=decoded_labels)
    return {"bleu": result_bleu["score"], "chrf": result_chrf["score"]}

### We now enter the fine-tuning phase of our model structure ###


#definition of seq2seq training arguments --- figure what these are/use case
args = Seq2SeqTrainingArguments(
    f"eng-nah-svo-translation",
    evaluation_strategy="no",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=64,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=3,
    predict_with_generate=True,
    fp16=False,
    push_to_hub=True,
)

#pass all information to trainer
trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

print("evaluate1")
print(trainer.evaluate(max_length=max_length))
print("trainer train 1")
trainer.train()
print("evaluate 2")
print(trainer.evaluate(max_length=max_length))
trainer.push_to_hub(tags="translation", commit_message="Training complete")
print("training model now")
model.train()


tokenized_datasets.set_format("torch")
train_dataloader = DataLoader(
    tokenized_datasets["train"],
    shuffle=True,
    collate_fn=data_collator,
    batch_size=8,
)
eval_dataloader = DataLoader(
    tokenized_datasets["test"], collate_fn=data_collator, batch_size=8, drop_last=True
)

model = AutoModelForSeq2SeqLM.from_pretrained("eng-nah-svo-translation")


optimizer = AdamW(model.parameters(), lr=2e-5)



accelerator = Accelerator()
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)



num_train_epochs = 3
num_update_steps_per_epoch = len(train_dataloader)
num_training_steps = num_train_epochs * num_update_steps_per_epoch

lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps,
)



model_name = "model"

output_dir = "./output"
repo = Repository("/mnt/storage/aatherton/hf_eng_fra_trans", clone_from="aatherton2024/hf_eng_fra_trans")


def postprocess(predictions, labels):
    predictions = predictions.cpu().numpy()
    labels = labels.cpu().numpy()

    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)

    # Replace -100 in the labels as we can't decode them.
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    # Some simple post-processing
    decoded_preds = [pred.strip() for pred in decoded_preds]
    decoded_labels = [[label.strip()] for label in decoded_labels]
    return decoded_preds, decoded_labels



# progress_bar = tqdm(range(num_training_steps))

# for epoch in range(num_train_epochs):
#     # Training
#     model.train()
#     for batch in train_dataloader:
#         outputs = model(**batch)
#         loss = outputs.loss
#         accelerator.backward(loss)

#         optimizer.step()
#         lr_scheduler.step()
#         optimizer.zero_grad()
#         progress_bar.update(1)

#     # Evaluation
#     model.eval()
#     for batch in tqdm(eval_dataloader):
#         with torch.no_grad():
#             generated_tokens = accelerator.unwrap_model(model).generate(
#                 batch["input_ids"],
#                 attention_mask=batch["attention_mask"],
#                 max_length=128,
#             )
#         labels = batch["labels"]

#         # Necessary to pad predictions and labels for being gathered
#         generated_tokens = accelerator.pad_across_processes(
#             generated_tokens, dim=1, pad_index=tokenizer.pad_token_id
#         )
#         labels = accelerator.pad_across_processes(labels, dim=1, pad_index=-100)

#         predictions_gathered = accelerator.gather(generated_tokens)
#         labels_gathered = accelerator.gather(labels)

#         decoded_preds, decoded_labels = postprocess(predictions_gathered, labels_gathered)
#         metric_bleu.add_batch(predictions=decoded_preds, references=decoded_labels)

#     results = metric_bleu.compute()
#     print(f"epoch {epoch}, BLEU score: {results['score']:.2f}")

#     # Save and upload
#     accelerator.wait_for_everyone()
#     unwrapped_model = accelerator.unwrap_model(model)
#     unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
#     if accelerator.is_main_process:
#         tokenizer.save_pretrained(output_dir)
#         repo.push_to_hub(
#             commit_message=f"Training in progress epoch {epoch}", blocking=False
#         )



# Replace this with your own checkpoint
model_checkpoint = "aatherton2024/eng-nah-svo-translation"
translator = pipeline("translation", model=model_checkpoint)
translator("Default to expanded threads")
print(translator(
    "you did not frichopize him"
))