notebooks/WECHSEL_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# WECHSEL Tutorial\n",
    "\n",
    "In this tutorial, we will use Langsfer to transfer a model trained in English to German with the [WECHSEL](https://arxiv.org/abs/2112.06598) method, similarily to one of the experiments described in the paper.\n",
    "\n",
    "WECHSEL is a cross-lingual language transfer method that efficiently initializes the embedding parameters of a language model in a target language using the embedding parameters from an existing model in a source language, facilitating more efficient training in the new language.\n",
    "\n",
    "The method requires as input:\n",
    "\n",
    "- a tokenizer in the source language,\n",
    "- a pre-trained language model in the source language,\n",
    "- a tokenizer in the target language,\n",
    "- 2 monolingual fastText embeddings for source and target languages respectively. \n",
    "  They can be obtained in one of 2 ways:\n",
    "    - using pre-trained fastText embeddings,\n",
    "    - trainining fastText embeddings from scratch.\n",
    "\n",
    "For the tutorial, we will use as much as possible the same parameters as described in the paper:\n",
    "\n",
    "- For the source model and tokenizer, we will use [roberta-base](https://huggingface.co/FacebookAI/roberta-base),\n",
    "- For the target tokenizer, we will train one from scratch,\n",
    "- For the fastText embeddings, we will download pre-trained models from [fastText's website](https://fasttext.cc/docs/en/crawl-vectors.html)\n",
    "\n",
    "For the sake of brevity, we will however use fewer training samples and steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "We begin by importing libraries and setting some defaults."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import warnings\n",
    "from typing import Generator\n",
    "\n",
    "import datasets\n",
    "import numpy as np\n",
    "import torch\n",
    "from transformers import (\n",
    "    AutoModel,\n",
    "    AutoTokenizer,\n",
    "    AutoModelForCausalLM,\n",
    "    DataCollatorForLanguageModeling,\n",
    "    TrainingArguments,\n",
    "    Trainer,\n",
    ")\n",
    "\n",
    "warnings.simplefilter(\"ignore\")\n",
    "\n",
    "# Constants\n",
    "SOURCE_MODEL_NAME = \"roberta-base\"\n",
    "DATASET_NAME = \"oscar-corpus/oscar\"\n",
    "DATASET_CONFIG_NAME = \"unshuffled_deduplicated_de\"\n",
    "DATASET_SIZE = 20000\n",
    "TRAIN_DATASET_SIZE = 16000\n",
    "TRAIN_BATCH_SIZE = 16\n",
    "GRADIENT_ACCUMULATION_STEPS = 8\n",
    "EVAL_STEPS = 4000\n",
    "MAX_TRAIN_STEPS = 48000\n",
    "LEARNING_RATE = 1e-4\n",
    "WEIGHT_DECAY = 0.01\n",
    "ADAM_EPSILON = 1e-6\n",
    "ADAM_BETA1 = 0.9\n",
    "ADAM_BETA2 = 0.98\n",
    "SEED = 16\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the following functions and classes from Langsfer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload\n",
    "from langsfer.high_level import wechsel\n",
    "from langsfer.embeddings import FastTextEmbeddings\n",
    "from langsfer.utils import download_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset\n",
    "\n",
    "We use the [datasets](https://huggingface.co/docs/datasets/index) library to load the [oscar](https://huggingface.co/datasets/oscar-corpus/oscar), which stands for **O**pen **S**uper-large **C**rawled **A**LMAnaCH co**R**pus, dataset's german configuration and then take a limited number of samples from it for training and validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasets.load_dataset(\n",
    "    DATASET_NAME,\n",
    "    DATASET_CONFIG_NAME,\n",
    "    split=\"train\",\n",
    "    streaming=True,\n",
    "    trust_remote_code=True,\n",
    ")\n",
    "dataset = dataset.shuffle(seed=SEED)\n",
    "dataset = dataset.take(DATASET_SIZE)\n",
    "train_dataset = dataset.take(TRAIN_DATASET_SIZE)\n",
    "val_dataset = dataset.skip(TRAIN_DATASET_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We take sample text from the validation set in order to compare tokenization between source and target tokenizers as well as for evaluating the generation of our trained model at the end. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mit Eva Mattes als Klara Blum. Eine Bäckerstochter stirbt in der Backröhre. Jetzt hat ihre Schwester (Julia Jentsch) Angst… Doppelbödig.\n",
      "in der rechten Armbeuge beim Öffnen des Mehlsilos zur Rettung der Bäckerstocher, welch ein Regiefehler! Der tiefergehende Sinn des Falles wird ansonsten auch nicht klar. Wirkt leider alles etwas zusammengeschustert.\n",
      "Wer spielte die Hauptrolle in Film \"The International\" und wurde als potenzieller James Bond-Nachfolger gehandelt?\n"
     ]
    }
   ],
   "source": [
    "sample_text = list(val_dataset.skip(10).take(1))[0][\"text\"]\n",
    "print(sample_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Embeddings and Tokenizers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load the source tokenizer as well as the source model and extract the input embeddings matrix from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaModel were not initialized from the model checkpoint at roberta-base and are newly initialized: ['roberta.pooler.dense.bias', 'roberta.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "source_tokenizer = AutoTokenizer.from_pretrained(SOURCE_MODEL_NAME)\n",
    "source_model = AutoModel.from_pretrained(SOURCE_MODEL_NAME)\n",
    "source_embeddings_matrix = source_model.get_input_embeddings().weight.detach().numpy()\n",
    "del source_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the source tokenizer to convert the sample text to tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of tokens 172, tokens: ['mit', 'ĠEva', 'ĠMatt', 'es', 'Ġal', 's', 'ĠKl', 'ara', 'ĠBl', 'um', '.', 'ĠE', 'ine', 'ĠB', 'Ã¤', 'cker', 'st', 'och', 'ter', 'Ġstir', 'bt', 'Ġin', 'Ġder', 'ĠBack', 'r', 'Ã¶', 'h', 're', '.', 'ĠJet', 'z', 't', 'Ġhat', 'Ġi', 'h', 're', 'ĠSchw', 'ester', 'Ġ(', 'Jul', 'ia', 'ĠJ', 'ents', 'ch', ')', 'ĠAng', 'st', 'âĢ¦', 'ĠDo', 'ppel', 'b', 'Ã¶', 'dig', '.', 'Ċ', 'in', 'Ġder', 'Ġre', 'ch', 'ten', 'ĠArm', 'be', 'uge', 'Ġbe', 'im', 'ĠÃĸ', 'ff', 'nen', 'Ġdes', 'ĠMeh', 'ls', 'il', 'os', 'Ġz', 'ur', 'ĠR', 'ett', 'ung', 'Ġder', 'ĠB', 'Ã¤', 'cker', 'st', 'oc', 'her', ',', 'Ġwel', 'ch', 'Ġe', 'in', 'ĠReg', 'ief', 'eh', 'ler', '!', 'ĠDer', 'Ġt', 'ief', 'er', 'ge', 'hend', 'e', 'ĠSinn', 'Ġdes', 'ĠFall', 'es', 'Ġw', 'ird', 'Ġan', 'son', 'sten', 'Ġa', 'uch', 'Ġn', 'icht', 'Ġk', 'lar', '.', 'ĠW', 'irk', 't', 'Ġle', 'ider', 'Ġall', 'es', 'Ġet', 'was', 'Ġz', 'us', 'amm', 'enges', 'ch', 'ust', 'ert', '.', 'Ċ', 'W', 'er', 'Ġsp', 'iel', 'te', 'Ġdie', 'ĠHau', 'pt', 'rol', 'le', 'Ġin', 'ĠFilm', 'Ġ\"', 'The', 'ĠInternational', '\"', 'Ġund', 'Ġw', 'urd', 'e', 'Ġal', 's', 'Ġpot', 'enzie', 'ller', 'ĠJames', 'ĠBond', '-', 'N', 'ach', 'fol', 'ger', 'Ġge', 'hand', 'elt', '?']\n"
     ]
    }
   ],
   "source": [
    "tokens = source_tokenizer.tokenize(sample_text)\n",
    "print(f\"Number of tokens {len(tokens)}, tokens: {tokens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We train a new target tokenizer using the same configuration as the source tokenizer using the training dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def batch_iterator(\n",
    "    dataset: datasets.Dataset, batch_size: int = 1000\n",
    ") -> Generator[str, None, None]:\n",
    "    for batch in dataset.iter(batch_size=batch_size):\n",
    "        yield batch[\"text\"]\n",
    "\n",
    "\n",
    "target_tokenizer = source_tokenizer.train_new_from_iterator(\n",
    "    batch_iterator(train_dataset), vocab_size=len(source_tokenizer)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then use the target tokenizer to convert the sample text to tokens and notice that the conversion creates fewer tokens than previously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of tokens 106, tokens: ['mit', 'ĠEva', 'ĠMatt', 'es', 'Ġals', 'ĠKl', 'ara', 'ĠBlum', '.', 'ĠEine', 'ĠBÃ¤cker', 'stochter', 'Ġstirbt', 'Ġin', 'Ġder', 'ĠBack', 'rÃ¶hre', '.', 'ĠJetzt', 'Ġhat', 'Ġihre', 'ĠSchwester', 'Ġ(', 'Julia', 'ĠJ', 'ent', 'sch', ')', 'ĠAngst', 'âĢ¦', 'ĠDoppel', 'bÃ¶', 'dig', '.', 'Ċ', 'in', 'Ġder', 'Ġrechten', 'ĠArmb', 'euge', 'Ġbeim', 'ĠÃĸffnen', 'Ġdes', 'ĠMehl', 'sil', 'os', 'Ġzur', 'ĠRettung', 'Ġder', 'ĠBÃ¤cker', 'st', 'ocher', ',', 'Ġwelch', 'Ġein', 'ĠReg', 'ief', 'ehler', '!', 'ĠDer', 'Ġtiefer', 'gehende', 'ĠSinn', 'Ġdes', 'ĠFall', 'es', 'Ġwird', 'Ġansonsten', 'Ġauch', 'Ġnicht', 'Ġklar', '.', 'ĠWir', 'kt', 'Ġleider', 'Ġalles', 'Ġetwas', 'Ġzusammen', 'gesch', 'uster', 't', '.', 'Ċ', 'Wer', 'Ġspielte', 'Ġdie', 'ĠHauptrolle', 'Ġin', 'ĠFilm', 'Ġ\"', 'The', 'ĠInternational', '\"', 'Ġund', 'Ġwurde', 'Ġals', 'Ġpoten', 'ziel', 'ler', 'ĠJames', 'ĠBond', '-', 'Nach', 'folger', 'Ġgehandelt', '?']\n"
     ]
    }
   ],
   "source": [
    "tokens = target_tokenizer.tokenize(sample_text)\n",
    "print(f\"Number of tokens {len(tokens)}, tokens: {tokens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then load pre-trained fasttext embeddings to use as auxiliary embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_auxiliary_embeddings = FastTextEmbeddings.from_model_name_or_path(\"en\")\n",
    "target_auxiliary_embeddings = FastTextEmbeddings.from_model_name_or_path(\"de\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After that, we download a bilinigual dictionary for English and German in order to be able to align the auxiliary embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "bilingual_dictionary_file = download_file(\n",
    "    \"https://raw.githubusercontent.com/CPJKU/wechsel/main/dicts/data/german.txt\",\n",
    "    \"/tmp/german.txt\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we open the file and read the first few lines, we can see that it maps English words to their German equivalent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'free': 'frei'},\n",
       " {'free': 'gratis'},\n",
       " {'free of charge': 'umsonst'},\n",
       " {'synonymous': 'synonym'},\n",
       " {'off': 'ab-'},\n",
       " {'suddenly': 'abrupt'},\n",
       " {'teetotal': 'abstinent'},\n",
       " {'pluralistic': 'plural'},\n",
       " {'house': 'Haus'},\n",
       " {'it': 'dat'}]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with bilingual_dictionary_file.open() as f:\n",
    "    dictionary_lines = [dict([f.readline().strip().split(\"\\t\")]) for _ in range(10)]\n",
    "\n",
    "dictionary_lines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We finally, instantiate the embedding initializer for WECHSEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_initializer = wechsel(\n",
    "    source_tokenizer=source_tokenizer,\n",
    "    source_embeddings_matrix=source_embeddings_matrix,\n",
    "    target_tokenizer=target_tokenizer,\n",
    "    target_auxiliary_embeddings=target_auxiliary_embeddings,\n",
    "    source_auxiliary_embeddings=source_auxiliary_embeddings,\n",
    "    bilingual_dictionary_file=bilingual_dictionary_file,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then initialize the target embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "63c7ebff3d3f4dbba88ceebd230d37ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Overlapping Tokens:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a850f025eff44e49a0a46103451c4b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Non-Overlapping Tokens: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "target_embeddings_matrix = embedding_initializer.initialize(seed=16, show_progress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have the initialized embeddings matrix, we can use it to replace the embeddings matrix in the source model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`\n"
     ]
    }
   ],
   "source": [
    "target_model_wechsel = AutoModelForCausalLM.from_pretrained(SOURCE_MODEL_NAME)\n",
    "\n",
    "# Resize its embedding layer\n",
    "target_model_wechsel.resize_token_embeddings(len(target_tokenizer))\n",
    "\n",
    "# Replace the source embeddings matrix with the target embeddings matrix\n",
    "target_model_wechsel.get_input_embeddings().weight.data = torch.as_tensor(\n",
    "    target_embeddings_matrix\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> We used `AutoModelForCausalLM` instead of `AutoModel` because we will train the newly initialized model for causal language modelling."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset preprocessing\n",
    "\n",
    "Before training, we must preprocess the training and validation sets by tokenizing the text, removing all other columns and then converting the resulting arrays to PyTorch tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.map(\n",
    "    lambda x: target_tokenizer(x[\"text\"], truncation=True),\n",
    "    batched=True,\n",
    "    remove_columns=dataset.column_names,\n",
    ")\n",
    "train_dataset = train_dataset.with_format(\"torch\")\n",
    "\n",
    "val_dataset = val_dataset.map(\n",
    "    lambda x: target_tokenizer(x[\"text\"], truncation=True),\n",
    "    batched=True,\n",
    "    remove_columns=dataset.column_names,\n",
    ")\n",
    "val_dataset = val_dataset.with_format(\"torch\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the training parameters and instantiate a [Trainer](https://huggingface.co/docs/transformers/en/main_classes/trainer) object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "max_steps is given, it will override any value given in num_train_epochs\n"
     ]
    }
   ],
   "source": [
    "data_collator = DataCollatorForLanguageModeling(tokenizer=target_tokenizer, mlm=False)\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"/tmp/wechsel\",\n",
    "    eval_strategy=\"steps\",\n",
    "    report_to=\"tensorboard\",\n",
    "    eval_steps=EVAL_STEPS // GRADIENT_ACCUMULATION_STEPS,\n",
    "    max_steps=MAX_TRAIN_STEPS // GRADIENT_ACCUMULATION_STEPS,\n",
    "    per_device_train_batch_size=TRAIN_BATCH_SIZE,\n",
    "    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,\n",
    "    learning_rate=LEARNING_RATE,\n",
    "    weight_decay=WEIGHT_DECAY,\n",
    "    adam_epsilon=ADAM_EPSILON,\n",
    "    adam_beta1=ADAM_BETA1,\n",
    "    adam_beta2=ADAM_BETA2,\n",
    "    bf16=True,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=target_model_wechsel,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=val_dataset,\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=target_tokenizer,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We evaluate the model before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluation loss before training: 13.195\n"
     ]
    }
   ],
   "source": [
    "eval_loss = trainer.evaluate()[\"eval_loss\"]\n",
    "print(f\"Evaluation loss before training: {eval_loss:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='6000' max='6000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [6000/6000 6:38:06, Epoch 95/9223372036854775807]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>0.705700</td>\n",
       "      <td>0.035726</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.015400</td>\n",
       "      <td>0.004960</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>0.006200</td>\n",
       "      <td>0.001831</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>0.004400</td>\n",
       "      <td>0.000643</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>0.001100</td>\n",
       "      <td>0.000399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3000</td>\n",
       "      <td>0.000800</td>\n",
       "      <td>0.000289</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3500</td>\n",
       "      <td>0.001400</td>\n",
       "      <td>0.000218</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4000</td>\n",
       "      <td>0.000500</td>\n",
       "      <td>0.000182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4500</td>\n",
       "      <td>0.000600</td>\n",
       "      <td>0.000150</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5000</td>\n",
       "      <td>0.000600</td>\n",
       "      <td>0.000134</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5500</td>\n",
       "      <td>0.000600</td>\n",
       "      <td>0.000124</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6000</td>\n",
       "      <td>0.000400</td>\n",
       "      <td>0.000120</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=6000, training_loss=0.061477883825699485, metrics={'train_runtime': 23892.0883, 'train_samples_per_second': 64.289, 'train_steps_per_second': 0.251, 'total_flos': 4.0437385940213856e+17, 'train_loss': 0.061477883825699485, 'epoch': 95.01041666666667})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We finally evaluate the model after the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluation loss after training: 0.000\n"
     ]
    }
   ],
   "source": [
    "eval_loss = trainer.evaluate()[\"eval_loss\"]\n",
    "print(f\"Evaluation loss after training: {eval_loss:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an additional evaluation, we take the sample text, truncate it and then make the trained model generate a completion for it "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original Text:\n",
      "mit Eva Mattes als Klara Blum. Eine Bäckerstochter stirbt in der Backröhre. Jetzt hat ihre Schwester (Julia Jentsch) Angst… Doppelbödig.\n",
      "in der rechten Armbeuge beim Öffnen des Mehlsilos zur Rettung der Bäckerstocher, welch ein Regiefehler! Der tiefergehende Sinn des Falles wird ansonsten auch nicht klar. Wirkt leider alles etwas zusammengeschustert.\n",
      "Wer spielte die Hauptrolle in Film \"The International\" und wurde als potenzieller James Bond-Nachfolger gehandelt?\n",
      "---\n",
      "Generated Text:\n",
      "<s>mit Eva Mattes als Klara Blum. Eine Bäckerstochter stirbt in der Backröhre. Jetzt hat ihre Schwester freundinschwesterOmamutter</s>\n"
     ]
    }
   ],
   "source": [
    "sample_input_ids = target_tokenizer(sample_text)[\"input_ids\"]\n",
    "shortened_input_ids = sample_input_ids[: len(sample_input_ids) // 3 - 13]\n",
    "\n",
    "generated_token_ids = (\n",
    "    trainer.model.generate(\n",
    "        torch.as_tensor(shortened_input_ids).reshape(1, -1).to(trainer.model.device),\n",
    "        max_length=300,\n",
    "        min_length=10,\n",
    "        top_p=0.9,\n",
    "        temperature=0.9,\n",
    "        repetition_penalty=2.0,\n",
    "    )\n",
    "    .detach()\n",
    "    .cpu()\n",
    "    .numpy()\n",
    "    .reshape(-1)\n",
    ")\n",
    "generated_tokens = target_tokenizer.decode(\n",
    "    generated_token_ids, add_special_tokens=False\n",
    ")\n",
    "print(\"Original Text:\")\n",
    "print(sample_text)\n",
    "print(\"---\")\n",
    "print(\"Generated Text:\")\n",
    "print(generated_tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generated text's quality is horrible and the model needs further training on more data,\n",
    "but this was just done for the sake of the tutorial and is not meant to be a full-blown model training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary\n",
    "\n",
    "In this tutorial, we have seen how to use WECHSEL in order to transfer a pre-trained language model to a new language."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langsfer-I0oHYpHZ-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}