training/flan_mt5_en_sp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-01-25 13:49:49.250254: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: SSE4.1 SSE4.2 AVX AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import torch\n",
    "import random\n",
    "import evaluate\n",
    "import transformers\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from typing import Optional\n",
    "from dataclasses import dataclass \n",
    "from time import perf_counter\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from datasets import load_dataset, disable_progress_bar\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSeq2SeqLM,\n",
    "    DataCollatorForSeq2Seq,\n",
    "    Seq2SeqTrainingArguments,\n",
    "    Seq2SeqTrainer,\n",
    "    EarlyStoppingCallback\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Config:\n",
    "    cache_dir: str = \"./data/english_to_spanish\" \n",
    "    data_dir: str = os.path.join(cache_dir, \"dataset\")\n",
    "    source_lang: str = \"en\"\n",
    "    target_lang: str = \"sp\"    \n",
    "    \n",
    "    batch_size: int = 16\n",
    "    num_workers: int = 4\n",
    "    seed: int = 42\n",
    "    max_source_length: int = 128\n",
    "    max_target_length: int = 128\n",
    "\n",
    "    lr: float = 0.0005\n",
    "    weight_decay: float = 0.01\n",
    "    epochs: int=20 #int = 20\n",
    "    device: torch.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    model_checkpoint: str = \"google/mt5-small\"\n",
    "\n",
    "    def __post_init__(self):\n",
    "        random.seed(self.seed)\n",
    "        np.random.seed(self.seed)\n",
    "        torch.manual_seed(self.seed)\n",
    "        torch.cuda.manual_seed_all(self.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "config=Config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset=pd.read_csv(os.path.join(config.data_dir, \"data.csv\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>english</th>\n",
       "      <th>spanish</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Go.</td>\n",
       "      <td>Ve.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Go.</td>\n",
       "      <td>Vete.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Go.</td>\n",
       "      <td>Vaya.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Go.</td>\n",
       "      <td>Váyase.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Hi.</td>\n",
       "      <td>Hola.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>118959</th>\n",
       "      <td>There are four main causes of alcohol-related ...</td>\n",
       "      <td>Hay cuatro causas principales de muertes relac...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>118960</th>\n",
       "      <td>There are mothers and fathers who will lie awa...</td>\n",
       "      <td>Hay madres y padres que se quedan despiertos d...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>118961</th>\n",
       "      <td>A carbon footprint is the amount of carbon dio...</td>\n",
       "      <td>Una huella de carbono es la cantidad de contam...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>118962</th>\n",
       "      <td>Since there are usually multiple websites on a...</td>\n",
       "      <td>Como suele haber varias páginas web sobre cual...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>118963</th>\n",
       "      <td>If you want to sound like a native speaker, yo...</td>\n",
       "      <td>Si quieres sonar como un hablante nativo, debe...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>118964 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  english  \\\n",
       "0                                                     Go.   \n",
       "1                                                     Go.   \n",
       "2                                                     Go.   \n",
       "3                                                     Go.   \n",
       "4                                                     Hi.   \n",
       "...                                                   ...   \n",
       "118959  There are four main causes of alcohol-related ...   \n",
       "118960  There are mothers and fathers who will lie awa...   \n",
       "118961  A carbon footprint is the amount of carbon dio...   \n",
       "118962  Since there are usually multiple websites on a...   \n",
       "118963  If you want to sound like a native speaker, yo...   \n",
       "\n",
       "                                                  spanish  \n",
       "0                                                     Ve.  \n",
       "1                                                   Vete.  \n",
       "2                                                   Vaya.  \n",
       "3                                                 Váyase.  \n",
       "4                                                   Hola.  \n",
       "...                                                   ...  \n",
       "118959  Hay cuatro causas principales de muertes relac...  \n",
       "118960  Hay madres y padres que se quedan despiertos d...  \n",
       "118961  Una huella de carbono es la cantidad de contam...  \n",
       "118962  Como suele haber varias páginas web sobre cual...  \n",
       "118963  Si quieres sonar como un hablante nativo, debe...  \n",
       "\n",
       "[118964 rows x 2 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset=dataset.sample(frac=1, random_state=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>english</th>\n",
       "      <th>spanish</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>26027</th>\n",
       "      <td>We need to chat soon.</td>\n",
       "      <td>Necesitamos charlar pronto.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33060</th>\n",
       "      <td>She married a musician.</td>\n",
       "      <td>Se casó con un músico.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7550</th>\n",
       "      <td>Do it right now.</td>\n",
       "      <td>Hazlo ahoritita.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37768</th>\n",
       "      <td>This is the ticket line.</td>\n",
       "      <td>Esta es la cola para sacar los billetes.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2075</th>\n",
       "      <td>No one came.</td>\n",
       "      <td>Nadie vino.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50057</th>\n",
       "      <td>The world is changing fast.</td>\n",
       "      <td>El mundo está cambiando rápidamente.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>98047</th>\n",
       "      <td>Tom did a good job organizing the workers.</td>\n",
       "      <td>Tom hizo un buen trabajo organizando a los tra...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5192</th>\n",
       "      <td>Tom left Mary.</td>\n",
       "      <td>Tom dejó a Mary.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>77708</th>\n",
       "      <td>They made me wait for a long time.</td>\n",
       "      <td>Me han hecho esperar mucho.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>98539</th>\n",
       "      <td>A short walk will bring you to the station.</td>\n",
       "      <td>Un paseo corto te llevará hasta la estación.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>118964 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           english  \\\n",
       "26027                        We need to chat soon.   \n",
       "33060                      She married a musician.   \n",
       "7550                              Do it right now.   \n",
       "37768                     This is the ticket line.   \n",
       "2075                                  No one came.   \n",
       "...                                            ...   \n",
       "50057                  The world is changing fast.   \n",
       "98047   Tom did a good job organizing the workers.   \n",
       "5192                                Tom left Mary.   \n",
       "77708           They made me wait for a long time.   \n",
       "98539  A short walk will bring you to the station.   \n",
       "\n",
       "                                                 spanish  \n",
       "26027                        Necesitamos charlar pronto.  \n",
       "33060                             Se casó con un músico.  \n",
       "7550                                    Hazlo ahoritita.  \n",
       "37768           Esta es la cola para sacar los billetes.  \n",
       "2075                                         Nadie vino.  \n",
       "...                                                  ...  \n",
       "50057               El mundo está cambiando rápidamente.  \n",
       "98047  Tom hizo un buen trabajo organizando a los tra...  \n",
       "5192                                    Tom dejó a Mary.  \n",
       "77708                        Me han hecho esperar mucho.  \n",
       "98539       Un paseo corto te llevará hasta la estación.  \n",
       "\n",
       "[118964 rows x 2 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset=dataset.head(10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset=dataset.reset_index()\n",
    "dataset=dataset.drop(columns='index')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>english</th>\n",
       "      <th>spanish</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>We need to chat soon.</td>\n",
       "      <td>Necesitamos charlar pronto.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>She married a musician.</td>\n",
       "      <td>Se casó con un músico.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Do it right now.</td>\n",
       "      <td>Hazlo ahoritita.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>This is the ticket line.</td>\n",
       "      <td>Esta es la cola para sacar los billetes.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>No one came.</td>\n",
       "      <td>Nadie vino.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9995</th>\n",
       "      <td>Did Tom tell you anything interesting?</td>\n",
       "      <td>¿Tom te dijo algo interesante?</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9996</th>\n",
       "      <td>Please keep this a secret.</td>\n",
       "      <td>Por favor, mantén esto en secreto.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9997</th>\n",
       "      <td>Have you lost your mind?</td>\n",
       "      <td>¿Has perdido el juicio?</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9998</th>\n",
       "      <td>He is unable to do it.</td>\n",
       "      <td>Él es incapaz de hacerlo.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9999</th>\n",
       "      <td>He lost the watch he bought yesterday.</td>\n",
       "      <td>Perdió el reloj que compró ayer.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>10000 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                     english  \\\n",
       "0                      We need to chat soon.   \n",
       "1                    She married a musician.   \n",
       "2                           Do it right now.   \n",
       "3                   This is the ticket line.   \n",
       "4                               No one came.   \n",
       "...                                      ...   \n",
       "9995  Did Tom tell you anything interesting?   \n",
       "9996              Please keep this a secret.   \n",
       "9997                Have you lost your mind?   \n",
       "9998                  He is unable to do it.   \n",
       "9999  He lost the watch he bought yesterday.   \n",
       "\n",
       "                                       spanish  \n",
       "0                  Necesitamos charlar pronto.  \n",
       "1                       Se casó con un músico.  \n",
       "2                             Hazlo ahoritita.  \n",
       "3     Esta es la cola para sacar los billetes.  \n",
       "4                                  Nadie vino.  \n",
       "...                                        ...  \n",
       "9995            ¿Tom te dijo algo interesante?  \n",
       "9996        Por favor, mantén esto en secreto.  \n",
       "9997                   ¿Has perdido el juicio?  \n",
       "9998                 Él es incapaz de hacerlo.  \n",
       "9999          Perdió el reloj que compró ayer.  \n",
       "\n",
       "[10000 rows x 2 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_validate_test_split(df, train_percent=.6, validate_percent=.2, seed=None):\n",
    "    np.random.seed(seed)\n",
    "    perm = np.random.permutation(df.index)\n",
    "    m = len(df.index)\n",
    "    train_end = int(train_percent * m)\n",
    "    validate_end = int(validate_percent * m) + train_end\n",
    "    train = df.iloc[perm[:train_end]]\n",
    "    validate = df.iloc[perm[train_end:validate_end]]\n",
    "    test = df.iloc[perm[validate_end:]]\n",
    "    return train, validate, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, validate, test = train_validate_test_split(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "train.to_csv(os.path.join(config.data_dir, \"train.csv\"), index=False, header=False)\n",
    "validate.to_csv(os.path.join(config.data_dir, \"validate.csv\"), index=False, header=False)\n",
    "test.to_csv(os.path.join(config.data_dir, \"test.csv\"), index=False, header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'train': './data/english_to_spanish/dataset/train.csv',\n",
       " 'validation': './data/english_to_spanish/dataset/validate.csv',\n",
       " 'test': './data/english_to_spanish/dataset/test.csv'}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_files = {\"train\": os.path.join(config.data_dir, \"train.csv\"),\n",
    "              \"validation\":os.path.join(config.data_dir, \"validate.csv\"), \n",
    "              \"test\": os.path.join(config.data_dir, \"test.csv\")}\n",
    "data_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb0df0b870834e5a8f2e30e3947346df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94759454b22b4e8f835d5d322a0440f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ee76656cdc34f70965041096831b82a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3625af1a7e6245bba2f78a6eda981a2d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating validation split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc7791c47bde47848dd42ad7e176f3ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['en', 'sp'],\n",
       "        num_rows: 6000\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['en', 'sp'],\n",
       "        num_rows: 2000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['en', 'sp'],\n",
       "        num_rows: 2000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_dict = load_dataset(\n",
    "    \"csv\",\n",
    "    delimiter=\",\",\n",
    "    column_names=[config.source_lang, config.target_lang],\n",
    "    data_files=data_files\n",
    ")\n",
    "dataset_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'en': 'He got all he wanted.', 'sp': 'Lograba todo lo que quería.'}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample = dataset_dict[\"train\"][0]\n",
    "sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting rouge_score\n",
      "  Downloading rouge_score-0.1.2.tar.gz (17 kB)\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting sacrebleu\n",
      "  Downloading sacrebleu-2.4.0-py3-none-any.whl.metadata (57 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.4/57.4 kB\u001b[0m \u001b[31m4.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: absl-py in /opt/conda/lib/python3.10/site-packages (from rouge_score) (2.0.0)\n",
      "Requirement already satisfied: nltk in /opt/conda/lib/python3.10/site-packages (from rouge_score) (3.8.1)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from rouge_score) (1.26.2)\n",
      "Requirement already satisfied: six>=1.14.0 in /opt/conda/lib/python3.10/site-packages (from rouge_score) (1.16.0)\n",
      "Collecting portalocker (from sacrebleu)\n",
      "  Downloading portalocker-2.8.2-py3-none-any.whl.metadata (8.5 kB)\n",
      "Requirement already satisfied: regex in /opt/conda/lib/python3.10/site-packages (from sacrebleu) (2023.10.3)\n",
      "Requirement already satisfied: tabulate>=0.8.9 in /opt/conda/lib/python3.10/site-packages (from sacrebleu) (0.9.0)\n",
      "Requirement already satisfied: colorama in /opt/conda/lib/python3.10/site-packages (from sacrebleu) (0.4.6)\n",
      "Collecting lxml (from sacrebleu)\n",
      "  Downloading lxml-5.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.5 kB)\n",
      "Requirement already satisfied: click in /opt/conda/lib/python3.10/site-packages (from nltk->rouge_score) (8.1.7)\n",
      "Requirement already satisfied: joblib in /opt/conda/lib/python3.10/site-packages (from nltk->rouge_score) (1.3.2)\n",
      "Requirement already satisfied: tqdm in /opt/conda/lib/python3.10/site-packages (from nltk->rouge_score) (4.66.1)\n",
      "Downloading sacrebleu-2.4.0-py3-none-any.whl (106 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m106.3/106.3 kB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading lxml-5.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (8.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.0/8.0 MB\u001b[0m \u001b[31m69.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m\n",
      "\u001b[?25hDownloading portalocker-2.8.2-py3-none-any.whl (17 kB)\n",
      "Building wheels for collected packages: rouge_score\n",
      "  Building wheel for rouge_score (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for rouge_score: filename=rouge_score-0.1.2-py3-none-any.whl size=24933 sha256=71b9af05d2fce11bb028e807e69c0824d89f5a2e0a44b706b73b8b96a428bf4f\n",
      "  Stored in directory: /home/sagemaker-user/.cache/pip/wheels/5f/dd/89/461065a73be61a532ff8599a28e9beef17985c9e9c31e541b4\n",
      "Successfully built rouge_score\n",
      "Installing collected packages: portalocker, lxml, sacrebleu, rouge_score\n",
      "Successfully installed lxml-5.1.0 portalocker-2.8.2 rouge_score-0.1.2 sacrebleu-2.4.0\n"
     ]
    }
   ],
   "source": [
    "!pip install rouge_score sacrebleu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eff4dac7033141438bb747f35036fb44",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/6.27k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4393c5469ddc4c2c9d7e17a03776f5b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/5.94k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "658dd6809d4a435a9e6cb09d5f14fe29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading extra modules:   0%|          | 0.00/1.55k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dcc3effa8b724729b714a4dfd6048c19",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading extra modules:   0%|          | 0.00/3.34k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1688c78dbaa5451c8e8133919e7a1fa7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/8.15k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rouge_score = evaluate.load(\"rouge\", cache_dir=config.cache_dir)\n",
    "bleu_score = evaluate.load(\"bleu\", cache_dir=config.cache_dir)\n",
    "sacrebleu_score = evaluate.load(\"sacrebleu\", cache_dir=config.cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbd83f0bbc0643f782d53ada842f5ea5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/82.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd5b567334924ef699c54872e64dd37d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/553 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e4d86f149294ce7adec074910fb3b5a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "spiece.model:   0%|          | 0.00/4.31M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "61b9dd7eeaf9444abd934c379d24929c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/99.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using the legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This means that tokens that come after special tokens will not be properly handled. We recommend you to read the related pull request available at https://github.com/huggingface/transformers/pull/24565\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/convert_slow_tokenizer.py:470: UserWarning: The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option which is not implemented in the fast tokenizers. In practice this means that the fast version of the tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these unknown tokens into a sequence of byte tokens matching the original piece of text.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(config.model_checkpoint, cache_dir=config.cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = config.model_checkpoint.split(\"/\")[-1]\n",
    "\n",
    "fine_tuned_model_checkpoint = os.path.join(\n",
    "    config.cache_dir,\n",
    "    f\"{model_name}_{config.source_lang}-{config.target_lang}\",\n",
    "    \"checkpoint-4500\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c3614519a1db4236acfd33a2fa303a60",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/1.20G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "06b5918c6158495e980c43d3d0be8092",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/147 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of parameters: 300176768\n"
     ]
    }
   ],
   "source": [
    "if os.path.isdir(fine_tuned_model_checkpoint):\n",
    "    do_train = False\n",
    "    model = AutoModelForSeq2SeqLM.from_pretrained(fine_tuned_model_checkpoint, cache_dir=config.cache_dir)\n",
    "else:\n",
    "    do_train = True\n",
    "    model = AutoModelForSeq2SeqLM.from_pretrained(config.model_checkpoint, cache_dir=config.cache_dir)\n",
    "\n",
    "print(\"number of parameters:\", model.num_parameters())\n",
    "#print(\"number of trainable parameters:\", model.print_trainable_parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_tokenize_fn(examples):\n",
    "    \"\"\"\n",
    "    Generate the input_ids and labels field for huggingface dataset/dataset dict.\n",
    "\n",
    "    Truncation is enabled where we cap the sentence to the max length. Padding will be done later\n",
    "    in a data collator, so we pad examples to the longest length within a mini-batch and not\n",
    "    the whole dataset.\n",
    "    \"\"\"\n",
    "    sources = examples[config.source_lang]\n",
    "    targets = examples[config.target_lang]\n",
    "    model_inputs = tokenizer(sources, max_length=config.max_source_length, truncation=True)\n",
    "\n",
    "    # setup the tokenizer for targets,\n",
    "    # huggingface expects the target tokenized ids to be stored in the labels field\n",
    "    # note, newer version of tokenizer supports a text_target argument, where we can create\n",
    "    # source and target sentences in one go\n",
    "    with tokenizer.as_target_tokenizer():\n",
    "        labels = tokenizer(targets, max_length=config.max_target_length, truncation=True)\n",
    "\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1d303b709ffa4d7f8a386f2736cc3eb7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/6000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:3635: UserWarning: `as_target_tokenizer` is deprecated and will be removed in v5 of Transformers. You can tokenize your labels by using the argument `text_target` of the regular `__call__` method (either in the same call as your input texts if you use the same keyword arguments, or in a separate call.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8f26cabe77c84b35b801e68183236479",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4b23b58bda24ebf9fb8b75782681233",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 6000\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 2000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 2000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_dict_tokenized = dataset_dict.map(\n",
    "    batch_tokenize_fn,\n",
    "    batched=True,\n",
    "    remove_columns=dataset_dict[\"train\"].column_names\n",
    ")\n",
    "dataset_dict_tokenized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = config.model_checkpoint.split(\"/\")[-1]\n",
    "output_dir = os.path.join(config.cache_dir, f\"{model_name}_{config.source_lang}-{config.target_lang}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Seq2SeqTrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    learning_rate=config.lr,\n",
    "    per_device_train_batch_size=config.batch_size,\n",
    "    per_device_eval_batch_size=config.batch_size,\n",
    "    weight_decay=config.weight_decay,\n",
    "    save_total_limit=2,\n",
    "    num_train_epochs=config.epochs,\n",
    "    predict_with_generate=True,\n",
    "    load_best_model_at_end=True,\n",
    "    greater_is_better=True,\n",
    "    metric_for_best_model=\"rougeL\",\n",
    "    gradient_accumulation_steps=8,\n",
    "    do_train=do_train,\n",
    "    #report_to=\"wandb\",  # enable logging to W&B\n",
    "    # careful when attempting to train t5 models on fp16 mixed precision,\n",
    "    # the model was trained on bfloat16 mixed precision, and mixing different mixed precision\n",
    "    # type might result in nan loss\n",
    "    # https://discuss.huggingface.co/t/mixed-precision-for-bfloat16-pretrained-models/5315\n",
    "    fp16=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(eval_pred):\n",
    "    \"\"\"\n",
    "    Compute rouge and bleu metrics for seq2seq model generated prediction.\n",
    "    \n",
    "    tip: we can run trainer.predict on our eval/test dataset to see what a sample\n",
    "    eval_pred object would look like when implementing custom compute metrics function\n",
    "    \"\"\"\n",
    "    predictions, labels = eval_pred\n",
    "    # Decode generated summaries, which is in ids into text\n",
    "    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
    "    # Replace -100 in the labels as we can't decode them\n",
    "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    # Decode labels, a.k.a. reference summaries into text\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "    result = rouge_score.compute(\n",
    "        predictions=decoded_preds,\n",
    "        references=decoded_labels,\n",
    "        rouge_types=[\"rouge1\", \"rouge2\", \"rougeL\"]\n",
    "    )\n",
    "    score = sacrebleu_score.compute(\n",
    "        predictions=decoded_preds,\n",
    "        references=decoded_labels\n",
    "    )\n",
    "    result[\"sacrebleu\"] = score[\"score\"]\n",
    "    return {k: round(v, 4) for k, v in result.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model,\n",
    "    args,\n",
    "    train_dataset=dataset_dict_tokenized[\"train\"],\n",
    "    eval_dataset=dataset_dict_tokenized[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics,\n",
    "    callbacks=[EarlyStoppingCallback()]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: You can find your API key in your browser here: https://wandb.ai/authorize\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Paste an API key from your profile and hit enter, or press ctrl+c to quit:"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  ········\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /home/sagemaker-user/.netrc\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.16.2 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.16.1"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/sagemaker-user/wandb/run-20240125_135113-0g0aaf7k</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/jorgev20/huggingface/runs/0g0aaf7k' target=\"_blank\">dazzling-wood-4</a></strong> to <a href='https://wandb.ai/jorgev20/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/jorgev20/huggingface' target=\"_blank\">https://wandb.ai/jorgev20/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/jorgev20/huggingface/runs/0g0aaf7k' target=\"_blank\">https://wandb.ai/jorgev20/huggingface/runs/0g0aaf7k</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a T5TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='920' max='920' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [920/920 3:49:15, Epoch 19/20]\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",
       "      <th>Rouge1</th>\n",
       "      <th>Rouge2</th>\n",
       "      <th>Rougel</th>\n",
       "      <th>Sacrebleu</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>2.892800</td>\n",
       "      <td>1.394342</td>\n",
       "      <td>0.556200</td>\n",
       "      <td>0.347500</td>\n",
       "      <td>0.544000</td>\n",
       "      <td>28.305300</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training elapsed time: 13814.898851252\n"
     ]
    }
   ],
   "source": [
    "# should take around 4117.78 seconds on a single V100 GPU\n",
    "if trainer.args.do_train:\n",
    "    os.environ[\"DISABLE_MLFLOW_INTEGRATION\"] = \"TRUE\"\n",
    "    t1_start = perf_counter()\n",
    "    train_output = trainer.train()\n",
    "    t1_stop = perf_counter()\n",
    "    print(\"Training elapsed time:\", t1_stop - t1_start)\n",
    "\n",
    "    # saving the model which allows us to leverage\n",
    "    # .from_pretrained(model_path)\n",
    "    trainer.save_model(fine_tuned_model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_translation(model, tokenizer, example):\n",
    "    \"\"\"print out the source, target and predicted raw text.\"\"\"\n",
    "    source = example[config.source_lang]\n",
    "    print(source)\n",
    "    target = example[config.target_lang]\n",
    "    print(target)\n",
    "    input_ids = tokenizer(source)[\"input_ids\"]\n",
    "    input_ids = torch.LongTensor(input_ids).view(1, -1).to(model.device)\n",
    "    print(input_ids)\n",
    "    #input_ids=input_ids[0]\n",
    "    generated_ids = model.generate(inputs=input_ids, max_length=20)\n",
    "    prediction = tokenizer.decode(generated_ids[0], skip_special_tokens=True)\n",
    "\n",
    "    print('source: ', source)\n",
    "    print('target: ', target)\n",
    "    print('prediction: ', prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'en': 'Compare this copy with the original.',\n",
       " 'sp': 'Compare usted esta copia con el original.'}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example = dataset_dict['validation'][0]\n",
    "example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compare this copy with the original.\n",
      "Compare usted esta copia con el original.\n",
      "tensor([[73343,   714, 27613,   514,   287,  4703,   260,     1]])\n",
      "source:  Compare this copy with the original.\n",
      "target:  Compare usted esta copia con el original.\n",
      "prediction:  Comparte este copia con el original.\n"
     ]
    }
   ],
   "source": [
    "generate_translation(model, tokenizer, example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_translation2(model, tokenizer, example):\n",
    "    \"\"\"print out the source, target and predicted raw text.\"\"\"\n",
    "    source = example\n",
    "    print(source)\n",
    "    #target = example[config.target_lang]\n",
    "    #print(target)\n",
    "    input_ids = tokenizer(source)[\"input_ids\"]\n",
    "    input_ids = torch.LongTensor(input_ids).view(1, -1).to(model.device)\n",
    "    print(input_ids)\n",
    "    #input_ids=input_ids[0]\n",
    "    generated_ids = model.generate(inputs=input_ids, max_length=20)\n",
    "    prediction = tokenizer.decode(generated_ids[0], skip_special_tokens=True)\n",
    "\n",
    "    print('source: ', source)\n",
    "    #print('target: ', target)\n",
    "    print('prediction: ', prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "example2='I enjoy my days with you'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I enjoy my days with you\n",
      "tensor([[ 336, 9070, 1037, 5382,  514,  521,    1]])\n",
      "source:  I enjoy my days with you\n",
      "prediction:  Adorcé mi días contigo.\n"
     ]
    }
   ],
   "source": [
    "generate_translation2(model, tokenizer, example2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}