{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10-minute self-cognition SFT\n",
    "\n",
    "Here is a demonstration of using python to perform self-cognition SFT of Qwen2.5-3B-Instruct. Through this tutorial, you can quickly understand some details of swift sft, which will be of great help in customizing ms-swift for you~\n",
    "\n",
    "Are you ready? Let's begin the journey...\n",
    "\n",
    "中文版：[魔搭教程](https://github.com/modelscope/modelscope-classroom/blob/main/LLM-tutorial/R.10%E5%88%86%E9%92%9F%E6%94%B9%E5%8F%98%E5%A4%A7%E6%A8%A1%E5%9E%8B%E8%87%AA%E6%88%91%E8%AE%A4%E7%9F%A5.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "# # install ms-swift\n",
    "# pip install ms-swift -U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import some libraries\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "from swift import get_model_processor, load_dataset, get_template\n",
    "from swift.dataset import EncodePreprocessor\n",
    "from swift.utils import get_logger, find_all_linears, get_model_parameter_info, plot_images, seed_everything\n",
    "from peft import get_peft_model, LoraConfig\n",
    "from swift.trainers import Seq2SeqTrainer, Seq2SeqTrainingArguments\n",
    "\n",
    "logger = get_logger()\n",
    "seed_everything(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters for training\n",
    "# model\n",
    "model_id_or_path = 'Qwen/Qwen2.5-3B-Instruct'  # model_id or model_path\n",
    "system = 'You are a helpful assistant.'\n",
    "output_dir = 'output'\n",
    "\n",
    "# dataset\n",
    "dataset = ['AI-ModelScope/alpaca-gpt4-data-zh#500', 'AI-ModelScope/alpaca-gpt4-data-en#500',\n",
    "           'swift/self-cognition#500']  # dataset_id or dataset_path\n",
    "data_seed = 42\n",
    "max_length = 2048\n",
    "split_dataset_ratio = 0.01  # Split validation set\n",
    "num_proc = 4  # The number of processes for data loading.\n",
    "# The following two parameters are used to override the placeholders in the self-cognition dataset.\n",
    "model_name = ['小黄', 'Xiao Huang']  # The Chinese name and English name of the model\n",
    "model_author = ['魔搭', 'ModelScope']  # The Chinese name and English name of the model author\n",
    "\n",
    "# lora\n",
    "lora_rank = 8\n",
    "lora_alpha = 32\n",
    "\n",
    "# training_args\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    learning_rate=1e-4,\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    gradient_checkpointing=True,\n",
    "    weight_decay=0.1,\n",
    "    lr_scheduler_type='cosine',\n",
    "    warmup_ratio=0.05,\n",
    "    report_to=['tensorboard'],\n",
    "    logging_first_step=True,\n",
    "    save_strategy='steps',\n",
    "    save_steps=50,\n",
    "    eval_strategy='steps',\n",
    "    eval_steps=50,\n",
    "    gradient_accumulation_steps=16,\n",
    "    num_train_epochs=1,\n",
    "    metric_for_best_model='loss',\n",
    "    save_total_limit=2,\n",
    "    logging_steps=5,\n",
    "    dataloader_num_workers=1,\n",
    "    data_seed=data_seed,\n",
    ")\n",
    "\n",
    "output_dir = os.path.abspath(os.path.expanduser(output_dir))\n",
    "logger.info(f'output_dir: {output_dir}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Obtain the model and template, and add a trainable Lora layer on the model.\n",
    "model, tokenizer = get_model_processor(model_id_or_path)\n",
    "logger.info(f'model_info: {model.model_info}')\n",
    "template = get_template(tokenizer, default_system=system, max_length=max_length)\n",
    "template.set_mode('train')\n",
    "\n",
    "target_modules = find_all_linears(model)\n",
    "lora_config = LoraConfig(task_type='CAUSAL_LM', r=lora_rank, lora_alpha=lora_alpha,\n",
    "                         target_modules=target_modules)\n",
    "model = get_peft_model(model, lora_config)\n",
    "logger.info(f'lora_config: {lora_config}')\n",
    "\n",
    "# Print model structure and trainable parameters.\n",
    "logger.info(f'model: {model}')\n",
    "model_parameter_info = get_model_parameter_info(model)\n",
    "logger.info(f'model_parameter_info: {model_parameter_info}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download and load the dataset, split it into a training set and a validation set,\n",
    "# and encode the text data into tokens.\n",
    "train_dataset, val_dataset = load_dataset(dataset, split_dataset_ratio=split_dataset_ratio, num_proc=num_proc,\n",
    "        model_name=model_name, model_author=model_author, seed=data_seed)\n",
    "\n",
    "logger.info(f'train_dataset: {train_dataset}')\n",
    "logger.info(f'val_dataset: {val_dataset}')\n",
    "logger.info(f'train_dataset[0]: {train_dataset[0]}')\n",
    "\n",
    "train_dataset = EncodePreprocessor(template=template)(train_dataset, num_proc=num_proc)\n",
    "val_dataset = EncodePreprocessor(template=template)(val_dataset, num_proc=num_proc)\n",
    "logger.info(f'encoded_train_dataset[0]: {train_dataset[0]}')\n",
    "\n",
    "# Print a sample\n",
    "template.print_inputs(train_dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the trainer and start the training.\n",
    "model.enable_input_require_grads()  # Compatible with gradient checkpointing\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    template=template,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=val_dataset,\n",
    ")\n",
    "trainer.train()\n",
    "\n",
    "last_model_checkpoint = trainer.state.last_model_checkpoint\n",
    "logger.info(f'last_model_checkpoint: {last_model_checkpoint}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the training loss.\n",
    "# You can also use the TensorBoard visualization interface during training by entering\n",
    "# `tensorboard --logdir '{output_dir}/runs'` at the command line.\n",
    "images_dir = os.path.join(output_dir, 'images')\n",
    "logger.info(f'images_dir: {images_dir}')\n",
    "plot_images(images_dir, training_args.logging_dir, ['train/loss'], 0.9)  # save images\n",
    "\n",
    "# Read and display the image.\n",
    "# The light yellow line represents the actual loss value,\n",
    "# while the yellow line represents the loss value smoothed with a smoothing factor of 0.9.\n",
    "from IPython.display import display\n",
    "from PIL import Image\n",
    "image = Image.open(os.path.join(images_dir, 'train_loss.png'))\n",
    "display(image)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hjt",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
