View file src/colab/hf_translation_fr_en_new.py - Download

# -*- coding: utf-8 -*-
"""hf_translation_fr_en_new.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1EIpgTHn6a79hCLuIwn9a7_dF5e9joojd

https://huggingface.co/docs/transformers/tasks/translation
"""

# !pip install transformers==4.28.0 datasets evaluate sacrebleu
!pip install transformers datasets evaluate sacrebleu

import torch

# Set device
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"device: {torch_device}")

"""Load OPUS Books dataset

"""

from datasets import load_dataset

books = load_dataset("opus_books", "en-fr")
books = books["train"].train_test_split(test_size=0.2)
print(books["train"][0])
print(books["test"][0])

"""Preprocess"""

from transformers import AutoTokenizer

checkpoint = "t5-small"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

source_lang = "fr"
target_lang = "en"
prefix = "translate French to English: "

def preprocess_function(examples):
    inputs = [prefix + example[source_lang] for example in examples["translation"]]
    targets = [example[target_lang] for example in examples["translation"]]
    model_inputs = tokenizer(inputs, text_target=targets, max_length=128, truncation=True)
    return model_inputs

tokenized_books = books.map(preprocess_function, batched=True)
print(tokenized_books)

from transformers import DataCollatorForSeq2Seq

data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=checkpoint)

"""Evaluate

"""

import evaluate

metric = evaluate.load("sacrebleu")

import numpy as np


def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [[label.strip()] for label in labels]

    return preds, labels


def compute_metrics(eval_preds):
    preds, labels = eval_preds
    if isinstance(preds, tuple):
        preds = preds[0]
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

    result = metric.compute(predictions=decoded_preds, references=decoded_labels)
    result = {"bleu": result["score"]}

    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
    result["gen_len"] = np.mean(prediction_lens)
    result = {k: round(v, 4) for k, v in result.items()}
    return result

"""Train"""

# !pip install accelerate==0.20.1
!pip install accelerate

from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer

model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)

training_args = Seq2SeqTrainingArguments(
    output_dir="my_awesome_opus_books_model",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=2,
    predict_with_generate=True,
    fp16=True,
    push_to_hub=False
)

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_books["train"],
    eval_dataset=tokenized_books["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# https://discuss.huggingface.co/t/how-to-turn-wandb-off-in-trainer/6237/10
import wandb
wandb.init(mode="disabled")

trainer.train()

"""Inference

"""

# text = "translate English to French: Legumes share resources with nitrogen-fixing bacteria."
text = "translate French to English: Les légumineuses partagent des ressources avec les bactéries fixatrices d'azote."
# from transformers import pipeline

# translator = pipeline("translation", model="my_awesome_opus_books_model")
# print(translator(text))


# from transformers import AutoTokenizer

# tokenizer = AutoTokenizer.from_pretrained("my_awesome_opus_books_model")
print("Tokenize...")
inputs = tokenizer(text, return_tensors="pt").input_ids.to(torch_device)

# from transformers import AutoModelForSeq2SeqLM

# model = AutoModelForSeq2SeqLM.from_pretrained("my_awesome_opus_books_model")
print("Convert to device...")
model1 = model.to(torch_device)
print("Translate...")
outputs = model1.generate(inputs, max_new_tokens=40, do_sample=True, top_k=30, top_p=0.95)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

from google.colab import drive
drive.mount("/content/drive")
torch.save(model1.state_dict(), "/content/drive/My Drive/hf_translation_fr_en.pth")

model2 = AutoModelForSeq2SeqLM.from_pretrained(checkpoint).to(torch_device)
model2.load_state_dict(torch.load("/content/drive/My Drive/hf_translation_fr_en.pth"))
outputs2 = model2.generate(inputs, max_new_tokens=40, do_sample=True, top_k=30, top_p=0.95)
print(tokenizer.decode(outputs2[0], skip_special_tokens=True))