{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/cu121\n",
    "!pip install transformers==4.37.0\n",
    "!pip install datasets==2.21.0\n",
    "!pip install accelerate==0.21.0\n",
    "!pip install rouge==1.0.1\n",
    "!pip install tqdm==4.66.5\n",
    "!pip install jieba==0.42.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "from transformers import AutoModelForSeq2SeqLM\n",
    "from transformers import DataCollatorForSeq2Seq\n",
    "from transformers import Seq2SeqTrainingArguments\n",
    "from transformers import Seq2SeqTrainer\n",
    "from datasets import load_dataset\n",
    "from rouge import Rouge\n",
    "import numpy as np\n",
    "import pickle\n",
    "import os\n",
    "import jieba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the dataset from Hugging Face Datasets\n",
    "# Note that we don't use to_list() here.\n",
    "# This is because we are going to use the map() function to process the dataset.\n",
    "data_name = \"hugcyp/LCSTS\"\n",
    "raw_train = load_dataset(data_name, split=\"train\")\n",
    "raw_valid = load_dataset(data_name, split=\"validation\")\n",
    "\n",
    "# To speed up evaluations during fine-tuning,\n",
    "# we only use a small subset of the validation set\n",
    "raw_small_valid = raw_valid.select(range(100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We use the multi-lingual T5 model for Chinese Abstractive Summarization.\n",
    "model_name = \"google/mt5-small\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "# Data Preprocessing takes time, so we save the preprocessed data with pickle.\n",
    "model_prefix = model_name.split(\"/\")[1]\n",
    "train_saved_pkl = f\"train_{model_prefix}.pkl\"\n",
    "valid_saved_pkl = f\"val_{model_prefix}.pkl\"\n",
    "small_val_saved_pkl = f\"val_{model_prefix}_100.pkl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We first check if the preprocessed data exists.\n",
    "# If the processed data exists, we load the data from the pickle file.\n",
    "if os.path.exists(small_val_saved_pkl):\n",
    "    assert os.path.exists(valid_saved_pkl)\n",
    "    with open(train_saved_pkl, \"rb\") as f:\n",
    "        train = pickle.load(f)\n",
    "    with open(valid_saved_pkl, \"rb\") as f:\n",
    "        valid = pickle.load(f)\n",
    "    with open(small_val_saved_pkl, \"rb\") as f:\n",
    "        small_valid = pickle.load(f)\n",
    "else:\n",
    "    token_replacement = [\n",
    "        [\"：\", \":\"],\n",
    "        [\"，\", \",\"],\n",
    "        [\"“\", '\"'],\n",
    "        [\"”\", '\"'],\n",
    "        [\"？\", \"?\"],\n",
    "        [\"……\", \"...\"],\n",
    "        [\"！\", \"!\"],\n",
    "    ]\n",
    "\n",
    "    def replace_tokens(examples):\n",
    "        # Substitute some punctuations to prevent too many [UNK] tokens\n",
    "        for k in [\"text\", \"summary\"]:\n",
    "            for i, _ in enumerate(examples[k]):\n",
    "                for tok in token_replacement:\n",
    "                    examples[k][i] = examples[k][i].replace(tok[0], tok[1])\n",
    "        return examples\n",
    "\n",
    "    def preprocess_function(examples):\n",
    "        examples = replace_tokens(examples)\n",
    "        model_inputs = tokenizer(examples[\"text\"], padding=True, truncation=True)\n",
    "        labels = tokenizer(\n",
    "            text_target=examples[\"summary\"],\n",
    "            max_length=200,\n",
    "            truncation=True,\n",
    "        )\n",
    "        model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "\n",
    "        return model_inputs\n",
    "\n",
    "    train = raw_train.map(preprocess_function, batched=True)\n",
    "    valid = raw_valid.map(preprocess_function, batched=True)\n",
    "    small_valid = raw_small_valid.map(preprocess_function, batched=True)\n",
    "\n",
    "    with open(train_saved_pkl, \"wb\") as f:\n",
    "        pickle.dump(train, f)\n",
    "    with open(valid_saved_pkl, \"wb\") as f:\n",
    "        pickle.dump(valid, f)\n",
    "    with open(small_val_saved_pkl, \"wb\") as f:\n",
    "        pickle.dump(small_valid, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the evaluation metric\n",
    "rouge_metric = Rouge()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    # Decode the predictions to sentences\n",
    "    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
    "    # We should first replace -100 in the labels with the pad token.\n",
    "    # -100 does not exist in the vocabulary, so we should restore it to the pad token.\n",
    "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    # Decode the labels to sentences\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "\n",
    "    # We use jieba to perform word-level evaluations with ROUGE\n",
    "    predictions = [\" \".join(jieba.lcut(o)) for o in decoded_preds]\n",
    "    references = [\" \".join(jieba.lcut(t)) for t in decoded_labels]\n",
    "    \n",
    "    # Set `avg=True` to compute the average scores for all samples\n",
    "    result = rouge_metric.get_scores(predictions, references, avg=True)\n",
    "    score = {f\"{rouge_i}_f\": v[\"f\"] for rouge_i, v in result.items()}\n",
    "    # Compute the average generation length\n",
    "    prediction_lens = [\n",
    "        np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions\n",
    "    ]\n",
    "    score[\"gen_len\"] = np.mean(prediction_lens)\n",
    "    return {k: v for k, v in score.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DataCollatorForSeq2Seq dynamically pads batched data and transforms padded labels into -100.\n",
    "# The operation provided by this object does a similar job like collate_fn.\n",
    "\n",
    "data_collator = DataCollatorForSeq2Seq(\n",
    "    tokenizer=tokenizer,\n",
    "    model=model_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use AutoModelForSeq2SeqLM for T5\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=\"./results/mt5\",\n",
    "    evaluation_strategy=\"steps\",\n",
    "    save_strategy=\"steps\",\n",
    "    eval_steps=1000,\n",
    "    save_steps=10000,\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=32,\n",
    "    per_device_eval_batch_size=32,\n",
    "    weight_decay=0.01,\n",
    "    save_total_limit=3,\n",
    "    num_train_epochs=3,\n",
    "    predict_with_generate=True,\n",
    "    logging_dir=f\"./logs/{model_prefix}\",\n",
    "    logging_steps=1,\n",
    "    push_to_hub=False,\n",
    ")\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train,\n",
    "    eval_dataset=small_valid, # We use the small validation set for evaluation\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "trainer.args._n_gpu = 1\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# After training, we evaluate the model on the validation set.\n",
    "results = trainer.predict(valid)\n",
    "for metric in [\"1\", \"2\", \"l\"]:\n",
    "    rouge_item = f\"test_rouge-{metric}\"\n",
    "    print(f\"{rouge_item}: \", results.metrics[rouge_item])"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}