1. Bi-MNTP Training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d09689f6-e9ba-482a-a3e6-31c2b05f696a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "037f7711-7760-422e-be4a-350fabf4d6b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "841d5810-bc64-497c-99c8-15400e698f9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mboyanov/miniconda3/lib/python3.11/site-packages/torch/cuda/__init__.py:141: UserWarning: CUDA initialization: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the available devices to be zero. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)\n",
      "  return torch._C._cuda_getDeviceCount() > 0\n"
     ]
    }
   ],
   "source": [
    "from bg2vec.arguments import *\n",
    "from bg2vec.data_util import *\n",
    "from bg2vec.model import *\n",
    "from bg2vec.training import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "de8866d5-bda1-43de-962f-521a20a598e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1a184357-2ecf-44c2-b230-3147a2ca9cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_args, data_args, training_args, custom_args = parser.parse_json_file(\"model_configurations/bggpt-7b.json\")\n",
    "\n",
    "config_kwargs = {\n",
    "    \"cache_dir\": model_args.cache_dir,\n",
    "    \"revision\": model_args.model_revision,\n",
    "    \"token\": model_args.token,\n",
    "    \"trust_remote_code\": model_args.trust_remote_code,\n",
    "}\n",
    "if training_args.gradient_checkpointing:\n",
    "    training_args.gradient_checkpointing_kwargs = {\"use_reentrant\": False}\n",
    "# Set seed before initializing model.\n",
    "\n",
    "transformers.set_seed(training_args.seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "347721df-c9f1-45df-b46b-636a0287b5f3",
   "metadata": {},
   "source": [
    "## Loading the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "904457cc-f6a3-4055-8601-2136a4768642",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import (\n",
    "    AutoConfig)\n",
    "\n",
    "config = AutoConfig.from_pretrained(\n",
    "    model_args.model_name_or_path, **config_kwargs\n",
    ")\n",
    "model_class = get_model_class(config)\n",
    "model_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe56443-e980-4a45-94f6-907fe925e49b",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_dtype = (\n",
    "    model_args.torch_dtype\n",
    "    if model_args.torch_dtype in [\"auto\", None]\n",
    "    else getattr(torch, model_args.torch_dtype)\n",
    ")\n",
    "model = model_class.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    from_tf=bool(\".ckpt\" in model_args.model_name_or_path),\n",
    "    config=config,\n",
    "    cache_dir=model_args.cache_dir,\n",
    "    revision=model_args.model_revision,\n",
    "    token=model_args.token,\n",
    "    trust_remote_code=model_args.trust_remote_code,\n",
    "    torch_dtype=torch_dtype,\n",
    "    low_cpu_mem_usage=model_args.low_cpu_mem_usage,\n",
    "    attn_implementation=model_args.attn_implementation,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "546f26f2-5b3f-4f65-8ddc-d081e8c4e7ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9825f73e-479a-4083-aff6-4cac79118cb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.model.layers[0].self_attn.is_causal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58d418ef-f54e-4ad7-a68f-dd4012c4d27e",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_model = initialize_peft(\n",
    "    model,\n",
    "    lora_r=custom_args.lora_r,\n",
    "    lora_alpha=2 * custom_args.lora_r,\n",
    "    lora_dropout=custom_args.lora_dropout,\n",
    ")\n",
    "model.set_model_for_peft(peft_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd998ac9-1fee-4f12-8c43-079e2d8e3c2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.model.base_model.model.layers[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c94bd869-c8b7-47b0-82bd-caa0f5440d82",
   "metadata": {},
   "source": [
    "## Setting up data collation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4a20063-96a6-4d1a-bbc0-74cb2da406b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "tokenizer_kwargs = {\n",
    "    \"cache_dir\": model_args.cache_dir,\n",
    "    \"use_fast\": model_args.use_fast_tokenizer,\n",
    "    \"revision\": model_args.model_revision,\n",
    "    \"token\": model_args.token,\n",
    "    \"trust_remote_code\": model_args.trust_remote_code,\n",
    "}\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "        model_args.model_name_or_path, **tokenizer_kwargs\n",
    ")\n",
    "tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4449374e-ea19-4898-873b-f3329f148efd",
   "metadata": {},
   "outputs": [],
   "source": [
    "if tokenizer.mask_token is None:\n",
    "    if custom_args.mask_token_type == \"blank\":\n",
    "        tokenizer.mask_token = \"_\"\n",
    "    elif custom_args.mask_token_type == \"eos\":\n",
    "        tokenizer.mask_token = tokenizer.eos_token\n",
    "    elif custom_args.mask_token_type == \"mask\":\n",
    "        tokenizer.add_tokens([\"<mask>\"])\n",
    "        tokenizer.mask_token = \"<mask>\"\n",
    "    else:\n",
    "        raise ValueError(\n",
    "            f\"mask_token_type {custom_args.mask_token_type} is not supported.\"\n",
    "        )\n",
    "\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "537b7c09-e41e-44b3-a8b5-e88180cb1ea0",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForLanguageModelingWithFullMasking(\n",
    "    tokenizer=tokenizer,\n",
    "    mlm_probability=data_args.mlm_probability\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "278f60c0-b741-4b98-888e-55ad2345d826",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2492a7f6-7da4-40f9-ae77-55af933e5a5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator.tokenizer.vocab['_']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b57f9c5d-57cb-43aa-a152-b281eed492ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator( (torch.randint(0, 10, (1, 10)), ))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4bd7760-b836-45a0-9df0-51d5a79ee194",
   "metadata": {},
   "source": [
    "## Loading up the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf6c3899-23aa-4053-85db-182ea59d36e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_datasets = datasets.load_from_disk(\"grouped_512\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c90ef94-5f74-4966-b2d2-753a1bfa11f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = tokenized_datasets[\"train\"]\n",
    "if data_args.max_train_samples is not None:\n",
    "    max_train_samples = min(len(train_dataset), data_args.max_train_samples)\n",
    "    train_dataset = train_dataset.select(range(max_train_samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bfae6f5-9f43-4754-86d9-6e7d2c162f64",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset = tokenized_datasets[\"validation\"]\n",
    "if data_args.max_eval_samples is not None:\n",
    "    max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)\n",
    "    eval_dataset = eval_dataset.select(range(max_eval_samples))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc331217-d0f3-4234-9b34-aa7e37f1c099",
   "metadata": {},
   "source": [
    "## Setting up the trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2fe5061f-96b7-4707-a63b-f76e425421ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import is_torch_tpu_available\n",
    "import evaluate\n",
    "from bg2vec.training import MNTPTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7ef36fc0-1a06-46b2-8e01-814b69aed3a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bg2vec.metrics import *\n",
    "\n",
    "evaluator = MetricEvaluator(model_args.cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74c2b9ee-7ad1-43c5-a31a-1b623348b0a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = MNTPTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=evaluator if training_args.do_eval and not is_torch_tpu_available()\n",
    "                              else None,\n",
    "    preprocess_logits_for_metrics=preprocess_logits_for_metrics\n",
    "    if training_args.do_eval and not is_torch_tpu_available()\n",
    "    else None,\n",
    "\n",
    ")\n",
    "\n",
    "trainer.add_callback(StopTrainingCallback(custom_args.stop_after_n_steps))\n",
    "trainer.callback_handler.remove_callback(transformers.integrations.integration_utils.WandbCallback)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81118d3c-535d-4b36-a36c-99ca09cdb303",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cffe07c5-2652-4641-8226-bc2113ae687a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_result = trainer.train()    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "470be2a4-c975-4cb2-b925-685f5c46ecd5",
   "metadata": {},
   "source": [
    "## Model export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f31d0534-def7-45b2-957b-e0fe4dd63555",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"bggpt-mntp-pretrained-2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f742637-a3df-46d0-8ec7-5ac9ddbaa9de",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model()  # Saves the tokenizer too for easy upload\n",
    "metrics = train_result.metrics\n",
    "\n",
    "max_train_samples = (\n",
    "    data_args.max_train_samples\n",
    "    if data_args.max_train_samples is not None\n",
    "    else len(train_dataset)\n",
    ")\n",
    "metrics[\"train_samples\"] = min(max_train_samples, len(train_dataset))\n",
    "\n",
    "trainer.log_metrics(\"train\", metrics)\n",
    "trainer.save_metrics(\"train\", metrics)\n",
    "trainer.save_state()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}