Translation Styling

Translation Styling

Image generated using the DALL-E model.

Introduction

Aiga went to a bookshop and bought a book about deep learning. She was very excited — at last she would properly refresh her knowledge before the final of the Olimpiada Sztucznej Inteligencji!

She came home, sat down comfortably with a mug of tea and started reading the first sentences:

"Zanurzenia zdaniowe otrzymane z modeli typu enkoder są zwykle lepsze od tylko dekoderowych."

"Uczenie ze wzmocnieniem z ludzkiego nadzoru typowo wykonuje się przez optymalizowanie polityki proksymalnej do modelu nagród."

[English gloss: the book translates every AI term into Polish. Word for word, the sentences read "Sentence embeddings obtained from encoder-type models are usually better than decoder-only ones." and "Reinforcement learning from human supervision is typically performed by optimising a policy proximal to the reward model."]

— What...? — muttered Aiga, looking at the pages in confusion.

She phoned a friend, who explained it to her straight away:

"Sentence embeddingi są lepsze, jeśli wyciągniesz je z encoderów niż z modeli decoder-only."

"RLHF typowo używa PPO, żeby optymalizować politykę względem reward modelu."

[English gloss: the friend keeps the English terms. The sentences mean "Sentence embeddings are better if you extract them from encoders than from decoder-only models." and "RLHF typically uses PPO to optimise the policy with respect to the reward model."]

Aiga breathed a sigh of relief. Now everything had become a little clearer.

The relief was only momentary, however — she still could not accept the style of the translations in the book. So she decided to train her own English-to-Polish translation model — one that would translate AI-related terms the way she likes!

Task

In this task, your goal is to adapt the translations generated by an existing machine translation model to a specific style – keeping the original form of certain industry terms from the field of machine learning.

The base model that we use is MarianMT – a model that translates from English into Polish, based on the encoder-decoder architecture. By default, this model translates all words, including specialist terminology. Your task is to implement a function that processes the input data and to fine-tune this model so that specialist terminology is not translated into Polish.

Data

The data available to you in this task are:

• Training set - contains English sentences, their translations into Polish and a list of keywords that should remain untranslated.

• Validation set - used to evaluate the effectiveness of your approach while training the model; it has the same format as the training set.

The examples have the following format:

{
    "keywords": ["explainable AI"],
    "en": "Developing explainable AI tools is crucial for trust in automated systems.",
    "pl": "Rozwijanie narzędzi explainable AI ma kluczowe znaczenie dla zaufania do zautomatyzowanych systemów.",
}

The test set on which your solution will ultimately be evaluated will not contain the Polish translations of the sentences.

Scoring Criterion

Your solution will be evaluated on hidden test data using the BLEU metric. The test set is similar to the validation set.

BLEU measures the overlap of n-grams (contiguous sequences of n adjacent words) between your translation and a single reference sentence – the greater the overlap, the higher the score.

Constraints

• Your solution will be tested on the Contest Platform without internet access and in an environment with a GPU.
• Training and evaluation of your solution on the Contest Platform may not take longer than 10 minutes.
• List of permitted libraries: torch, pandas, numpy, nltk, transformers, datasets, matplotlib.

Submission Files

• This notebook, completed with your solution.

Evaluation

Remember that during checking, the FINAL_EVALUATION_MODE flag will be set to True.

For this task you can obtain between 0 and 100 points. If your model achieves a BLEU score lower than 0.82 on the (secret) test set on the Contest Platform, your solution will receive 0 points. If it achieves a score higher than 0.86, you will receive the maximum number of points. Between these thresholds, your score will scale linearly and will be rounded to the nearest integer.

$$\mathrm{points} = \begin{cases} 100 & \text{if } \mathrm{BLEU} > 0.86 \\ 0 & \text{if } \mathrm{BLEU} < 0.82 \\ 100 \cdot \dfrac{\mathrm{BLEU} - 0.82}{0.86 - 0.82} & \text{otherwise} \end{cases}$$

Starter Code

In this section we initialise the environment by importing the required libraries and functions. The prepared code will make it easier for you to work with the data efficiently and to build the proper solution.

######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

FINAL_EVALUATION_MODE = False  # When your solution is checked, we will change this value to True

import json
import os
import torch

import matplotlib.pyplot as plt
import numpy as np
from datasets import Dataset
from nltk.translate.bleu_score import sentence_bleu
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

from transformers import set_seed
import random

seed = 42

os.environ["PYTHONHASHSEED"] = str(seed)

random.seed(seed)
np.random.seed(seed)

torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

set_seed(seed)

######################### DO NOT CHANGE THIS CELL ##########################

def load_dataset(json_path: str) -> Dataset:
    with open(json_path, "r", encoding="utf-8") as f:
        lines = f.readlines()

    dataset = []
    for line in lines:
        item = json.loads(line)
        dataset.append(
            {
                "en": item["translation"]["en"],
                "pl": item["translation"]["pl"],
                "keywords": ",".join(item.get("keywords", []))
            }
        )

    return Dataset.from_list(dataset)
    

train_dataset = load_dataset("train_dataset.jsonl")
print(f"Loaded {len(train_dataset)} training examples.")

val_dataset = load_dataset("valid_dataset.jsonl")
print(f"Loaded {len(val_dataset)} validation examples.")

######################### DO NOT CHANGE THIS CELL ##########################

def evaluate_model(model, tokenizer, dataset, process_example_fn, batch_size=64, device="cuda", verbose=True):
    """
    Function that evaluates the model using the BLEU metric.
    Arguments:
        model: The model to evaluate.
        tokenizer: The tokenizer used to process the text.
        dataset: The dataset for evaluation.
        process_example_fn: Function that processes the examples:
            (en_sentence: str, keywords: List[str]) -> (input_sentence: str).
        verbose: Whether to display details for the first few examples.
    Returns:
        float: The mean BLEU score for the dataset.
    """
    model.eval()
    model.to(device)

    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)

    bleu_scores = []

    for batch_idx, batch in enumerate(dataloader):
        orig_sentences = batch["en"]
        reference_translations = batch["pl"]
        keywords = batch["keywords"]

        model_inputs = [
            process_example_fn(orig, kws) 
            for orig, kws in zip(orig_sentences, keywords)
        ]

        inputs = tokenizer(model_inputs, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)

        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=64)

        hypotheses = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        for i, (ref, hyp) in enumerate(zip(reference_translations, hypotheses)):
            score = sentence_bleu([ref], hyp)
            bleu_scores.append(score)

            if verbose and batch_idx * batch_size + i < 5:
                print(f"Example {batch_idx * batch_size + i}:")
                print(f"Original: {orig_sentences[i]}")
                print(f"Processed: {model_inputs[i]}")
                print(f"Reference: {ref}")
                print(f"Hypothesis: {hyp}")
                print(f"BLEU score: {score:.4f}")
                print("-" * 10)

    bleu_score = sum(bleu_scores) / len(bleu_scores)
    return bleu_score


def compute_score(bleu_score: float) -> float:
    """
    Computes the points score based on the value of the BLEU metric.
    """
    lower_bound = 0.82
    upper_bound = 0.86

    if bleu_score <= lower_bound:
        return 0
    elif lower_bound < bleu_score < upper_bound:
        return int(round(100 * (bleu_score - lower_bound) / (upper_bound - lower_bound)))
    else:
        return 100

def process_example(en: str, keywords: str) -> str:
    """
    Function that converts the evaluation examples into the input text
    for the model.
    During training you may use this function, but you do not have to.
    Arguments:
        en: Text in English.
        keywords: Keywords separated by commas.
    """
    # TODO: Implement the function

    return en

# TODO: Model training.
# You are not allowed to change the tokenizer.
# Remember that you have the "gsarti/opus-mt-tc-en-pl" model and tokenizer at your disposal.

tokenizer = AutoTokenizer.from_pretrained("gsarti/opus-mt-tc-en-pl")
model = AutoModelForSeq2SeqLM.from_pretrained("gsarti/opus-mt-tc-en-pl")

# TODO: ...

my_model = model  # TODO: You must assign your final model to the variable "my_model".

######################### DO NOT CHANGE THIS CELL ##########################

if not FINAL_EVALUATION_MODE:
    tokenizer = AutoTokenizer.from_pretrained("gsarti/opus-mt-tc-en-pl")
    bleu_score = evaluate_model(
        model=my_model, 
        tokenizer=tokenizer, 
        dataset=val_dataset,
        process_example_fn=process_example,
    )
    print(f"BLEU score on the validation set: {bleu_score:.4f}")

    score = compute_score(bleu_score)
    print(f"Points on the validation set: {score}")

######################### DO NOT CHANGE THIS CELL ##########################

if FINAL_EVALUATION_MODE:
    import cloudpickle

    OUTPUT_PATH = "file_output"

    FUNCTION_OUTPUT_PATH = os.path.join(OUTPUT_PATH, "your_function.pkl")
    MODEL_OUTPUT_PATH = os.path.join(OUTPUT_PATH, "your_model.pkl")

    if not os.path.exists(OUTPUT_PATH):
        os.makedirs(OUTPUT_PATH)

    with open(FUNCTION_OUTPUT_PATH, "wb") as f:
        cloudpickle.dump(process_example, f)

    with open(MODEL_OUTPUT_PATH, "wb") as f:
        cloudpickle.dump(my_model, f)