note.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed2fb84",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers\n",
    "!pip install --upgrade ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "703616bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from json import loads, load\n",
    "from random import randint\n",
    "\n",
    "from torch.utils.data import Dataset\n",
    "from torch import tensor, cuda, randint as torch_rand\n",
    "\n",
    "from transformers import TrainingArguments, AutoConfig, RobertaConfig, GPT2Config, GPTJConfig\n",
    "from transformers import RobertaForMaskedLM, GPT2LMHeadModel, GPTJModel\n",
    "from transformers import RobertaTokenizerFast, AutoTokenizer, DataCollatorForLanguageModeling\n",
    "from transformers import DefaultFlowCallback\n",
    "from transformers.trainer_callback import TrainerState, TrainerControl, TrainingArguments, IntervalStrategy\n",
    "from transformers import pipeline, Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9df3a5cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "for c in range(0,cuda.device_count()):\n",
    "    print(cuda.get_device_name(c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c66024d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class JsonDataset(Dataset):\n",
    "    def __init__(self, jpath):\n",
    "        if isinstance(jpath, str):\n",
    "            with open(jpath, \"r\", encoding=\"utf-8\") as jf:\n",
    "                self.examples = list(jf)\n",
    "        else:\n",
    "            self.examples = []\n",
    "            for jp in jpath:\n",
    "                with open(jp, \"r\", encoding=\"utf-8\") as jf:\n",
    "                    self.examples += list(jf)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.examples)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return tensor(loads(self.examples[i])).long()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd17fd33",
   "metadata": {},
   "outputs": [],
   "source": [
    "false = False\n",
    "true = True\n",
    "\n",
    "config = {\n",
    "    \"paths\": {\n",
    "        \"main_path\": \"C:/Users/Administrator/Desktop/training\",\n",
    "        \"train_path\": \"%main_path%/mini_train.jsonl\",\n",
    "        \"dev_path\": \"%main_path%/mini_dev.jsonl\",\n",
    "        \"tokenizer_path\": \"%main_path%/tokenizer.json\",\n",
    "        \"model_folder\": \"saved\",\n",
    "        \"pretrained\": \"\"\n",
    "    },\n",
    "\n",
    "    \"model-options\": {\n",
    "        \"model_type\": \"roberta-base\",\n",
    "        \"resume-from-checkpoint\": false,\n",
    "        \"output_from_model\": true\n",
    "    },\n",
    "\n",
    "    \"training-options\": {\n",
    "        \"num_train_epochs\": 3,\n",
    "        \"per_device_train_batch_size\": 2,\n",
    "        \"per_device_eval_batch_size\": 2,\n",
    "        \"learning_rate\": 0.00008,\n",
    "        \"weight_decay\": 0.1,\n",
    "        \"warmup_steps\": 2000,\n",
    "\n",
    "        \"save_steps\": 50000,\n",
    "        \"eval_steps\": 50000,\n",
    "        \"save_total_limit\": 1,\n",
    "        \"load_best_model_at_end\": false,\n",
    "        \"overwrite_output_dir\": true,\n",
    "        \"evaluation_strategy\": \"epoch\"\n",
    "    },\n",
    "\n",
    "    \"misc\": {\n",
    "        \"encoded_file_keyword\": \"_encoded_\",\n",
    "        \"default_gen_input\": \"\"\n",
    "    }\n",
    "}\n",
    "\n",
    "model_params = {\n",
    "    \"num_attention_heads\": 12,\n",
    "    \"num_hidden_layers\": 12,\n",
    "    \"hidden_size\": 768,\n",
    "\n",
    "    \"max_position_embeddings\": 514,\n",
    "    \"type_vocab_size\": 1,\n",
    "    \"attention_probs_dropout_prob\": 0.1,\n",
    "    \"hidden_act\": \"gelu\",\n",
    "    \"hidden_dropout_prob\": 0.1,\n",
    "    \"initializer_range\": 0.02,\n",
    "    \"intermediate_size\": 4096,\n",
    "    \"layer_norm_eps\": 0.00001\n",
    "}\n",
    "\n",
    "examples = [\n",
    "    \"Ana ide u <mask>.\",\n",
    "    \"Osnovna <mask> Vuk Karadžić\",\n",
    "    \"Kupio sam dva <mask> i mleko.\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82c20009",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model(model_type, fast_tokenizer, pretrained=\"\", model_params=None):\n",
    "    if pretrained:\n",
    "        if \"roberta\" in model_type:\n",
    "            return RobertaForMaskedLM.from_pretrained(pretrained)\n",
    "        elif \"gpt2\" in model_type:\n",
    "            return GPT2LMHeadModel.from_pretrained(pretrained)\n",
    "        elif \"gptj\" in model_type:\n",
    "            return GPTJModel.from_pretrained(pretrained)\n",
    "\n",
    "    else:\n",
    "        if not model_params:\n",
    "            with open(\"training-configs/\" + model_type + \".json\", \"r\") as mf:\n",
    "                model_params = load(mf)\n",
    "        return create_model(model_type, fast_tokenizer, model_params)\n",
    "\n",
    "\n",
    "def create_model(model_type, fast_tokenizer, model_params):\n",
    "    if \"roberta\" in model_type:\n",
    "        model_config = RobertaConfig(**model_params)\n",
    "    elif \"gpt2\" in model_type:\n",
    "        model_config = GPT2Config(**model_params)\n",
    "    elif \"gptj\" in model_type:\n",
    "        model_config = GPTJConfig(**model_params)\n",
    "\n",
    "    else:\n",
    "        model_config = AutoConfig()\n",
    "\n",
    "    model_config.vocab_size = fast_tokenizer.vocab_size\n",
    "    model_config.bos_token_id = fast_tokenizer.bos_token_id\n",
    "    model_config.eos_token_id = fast_tokenizer.bos_token_id\n",
    "\n",
    "    if \"roberta\" in model_type:\n",
    "        return RobertaForMaskedLM(config=model_config)\n",
    "    elif \"gpt2\" in model_type:\n",
    "        return GPT2LMHeadModel(config=model_config)\n",
    "    elif \"gptj\" in model_type:\n",
    "        return GPTJModel(config=model_config)\n",
    "\n",
    "\n",
    "def load_tokenizer(model_type, tokenizer_path):\n",
    "    if \"roberta\" in model_type:\n",
    "        return RobertaTokenizerFast(tokenizer_file=tokenizer_path,\n",
    "                                    pad_token=\"<pad>\", unk_token=\"<unk>\", mask_token=\"<mask>\")\n",
    "    elif \"gpt\" in model_type:\n",
    "        return RobertaTokenizerFast(tokenizer_file=tokenizer_path, padding=False, pad_token=\"<pad>\")\n",
    "    else:\n",
    "        return AutoTokenizer()\n",
    "\n",
    "\n",
    "def collator(model_type, fast_tokenizer):\n",
    "    if \"roberta\" in model_type:\n",
    "        return DataCollatorForLanguageModeling(\n",
    "            mlm=True,\n",
    "            mlm_probability=0.15,\n",
    "            tokenizer=fast_tokenizer,\n",
    "        )\n",
    "    elif \"gpt\" in model_type:\n",
    "        return DataCollatorForLanguageModeling(\n",
    "            tokenizer=fast_tokenizer,\n",
    "            mlm=False,\n",
    "        )\n",
    "    else:\n",
    "        return DataCollatorForLanguageModeling(\n",
    "            tokenizer=fast_tokenizer\n",
    "        )\n",
    "\n",
    "\n",
    "def load_configs(cfg=None, cfgpath=\"training-configs/config.json\"):\n",
    "    if not cfg:\n",
    "        with open(cfgpath, \"r\") as cf:\n",
    "            cfg = load(cf)\n",
    "\n",
    "    # paths\n",
    "    main_path = cfg[\"paths\"][\"main_path\"]\n",
    "    newpaths = {x: process_path(y, \"%main_path%\", main_path) for (x, y) in cfg[\"paths\"].items()}\n",
    "\n",
    "    # model and training parameters\n",
    "    options = cfg[\"model-options\"]\n",
    "\n",
    "    training_options = cfg[\"training-options\"]\n",
    "    training_options[\"output_dir\"] = newpaths[\"model_folder\"]\n",
    "    training_options[\"remove_unused_columns\"] = False\n",
    "\n",
    "    # Training args fill\n",
    "    args = TrainingArguments(**training_options)\n",
    "    efk = cfg[\"misc\"][\"encoded_file_keyword\"]\n",
    "    default_input = cfg[\"misc\"][\"default_gen_input\"]\n",
    "    return newpaths, options, args, efk, default_input\n",
    "\n",
    "\n",
    "def process_path(path, key, replace_path):\n",
    "    if isinstance(path, str):\n",
    "        return path.replace(key, replace_path)\n",
    "    else:\n",
    "        results = []\n",
    "        for x in path:\n",
    "            results.append(x.replace(key, replace_path))\n",
    "        return results\n",
    "\n",
    "\n",
    "def get_examples(examples=None, examples_path=\"training-configs/fill_mask_examples.json\"):\n",
    "    if not examples:\n",
    "        with open(examples_path, \"r\") as ef:\n",
    "            examples = load(ef)\n",
    "    return examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49bdfed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "paths, model_options, training_args, encoded_file_keyword, default_gen_input = load_configs(config)\n",
    "fill_test_examples = get_examples(examples)\n",
    "tokenizer = load_tokenizer(model_options[\"model_type\"], paths[\"tokenizer_path\"])\n",
    "data_collator = collator(model_options[\"model_type\"], tokenizer)\n",
    "model = get_model(model_options[\"model_type\"], tokenizer, paths[\"pretrained\"], model_params)\n",
    "device = \"cuda:0\" if cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f49e821",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fill_examples(mod, tok):\n",
    "    # Create a Fill mask pipeline\n",
    "    fill_mask = pipeline(\n",
    "        \"fill-mask\",\n",
    "        model=mod,\n",
    "        tokenizer=tok,\n",
    "        device=device,\n",
    "        top_k=3\n",
    "    )\n",
    "    examples = []\n",
    "    for example in fill_test_examples:\n",
    "        examples.append([x[\"sequence\"] for x in fill_mask(example)])\n",
    "    return examples\n",
    "\n",
    "\n",
    "def generate(model, context, length=20, temperature=0.75):\n",
    "    encoded_input = context.to(device)\n",
    "    output = model.generate(\n",
    "        **encoded_input,\n",
    "        bos_token_id=randint(1, 50000),\n",
    "        do_sample=True,\n",
    "        top_k=0,\n",
    "        max_length=length,\n",
    "        temperature=temperature,\n",
    "        no_repeat_ngram_size=3,\n",
    "        # top_p=0.95,\n",
    "        num_return_sequences=1,\n",
    "        pad_token_id=0\n",
    "    )\n",
    "\n",
    "    return output\n",
    "\n",
    "\n",
    "def generatetion_test(mod, tok, samples=3, length=24, context=default_gen_input, temp=0.75):\n",
    "    outs = []\n",
    "    if context == \"\":\n",
    "        tokens = torch_rand(low=260, high=52000, size=(1,))\n",
    "        context = tok.decode(tokens, skip_special_tokens=True)\n",
    "\n",
    "    context = tok(context, return_tensors=\"pt\")\n",
    "    cl = context.data[\"input_ids\"].size()[1]\n",
    "\n",
    "    for x in range(samples):\n",
    "        output = generate(mod, context=context, length=length + cl, temperature=temp)\n",
    "\n",
    "        decoded_output = []\n",
    "        for sample in output:\n",
    "            sample = sample[cl:]\n",
    "            decoded_output.append(tokenizer.decode(sample, skip_special_tokens=True))\n",
    "\n",
    "        outs.append(\"\".join(decoded_output))\n",
    "\n",
    "    return outs\n",
    "\n",
    "\n",
    "def test(mod, tok=tokenizer):\n",
    "    if \"roberta\" in model_options[\"model_type\"]:\n",
    "        return fill_examples(mod, tok)\n",
    "    elif \"gpt\" in model_options[\"model_type\"]:\n",
    "        return generatetion_test(mod, tok)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6821d608",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDefaultFlowCallback(DefaultFlowCallback):\n",
    "    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):\n",
    "        # Log\n",
    "        if state.global_step == 1 and args.logging_first_step:\n",
    "            control.should_log = True\n",
    "        if args.logging_strategy == IntervalStrategy.STEPS and state.global_step % args.logging_steps == 0:\n",
    "            control.should_log = True\n",
    "\n",
    "        # Evaluate\n",
    "        if (\n",
    "                args.evaluation_strategy == IntervalStrategy.STEPS\n",
    "                and state.global_step % args.eval_steps == 0\n",
    "                and args.eval_delay <= state.global_step\n",
    "        ):\n",
    "            control.should_evaluate = True\n",
    "\n",
    "        # Save\n",
    "        if (\n",
    "                args.save_strategy == IntervalStrategy.STEPS\n",
    "                and args.save_steps > 0\n",
    "                and state.global_step % args.save_steps == 0\n",
    "        ):\n",
    "            control.should_save = True\n",
    "            examples = test(kwargs[\"model\"])\n",
    "            examples = [e for ee in examples for e in ee]\n",
    "            with open(paths[\"model_folder\"] + \"/experiments.log\", \"a+\", encoding=\"utf-8\") as lf:\n",
    "                lf.write(\"\\t\".join(examples))\n",
    "                lf.write(\"\\n\")\n",
    "\n",
    "        # End training\n",
    "        if state.global_step >= state.max_steps:\n",
    "            control.should_training_stop = True\n",
    "\n",
    "        return control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebde8fb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset = JsonDataset(paths[\"dev_path\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b85b7ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = JsonDataset(paths[\"train_path\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e7ea3be",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    #prediction_loss_only=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af0baf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_options[\"output_from_model\"]:\n",
    "    trainer.remove_callback(DefaultFlowCallback)\n",
    "    trainer.add_callback(CustomDefaultFlowCallback)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c02d7669",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train(resume_from_checkpoint=model_options[\"resume-from-checkpoint\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9790a856",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}