emoji-lm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, transformers, pyreft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_examples = [\n",
    "    [\"What are common symptoms of a cold?\", \"🤧😷🤒\"],\n",
    "    [\"How do you treat a headache?\", \"💊🧴💆‍♂️\"],\n",
    "    [\"What are the side effects of aspirin?\", \"💊🤕😴\"],\n",
    "    [\"How do you measure blood pressure?\", \"🩺💉📈\"],\n",
    "    [\"What should I do if I have a fever?\", \"🌡️💊🛏️\"],\n",
    "    [\"What are the signs of a heart attack?\", \"❤️🚨💔\"],\n",
    "    [\"How often should I get a health check-up?\", \"🗓️🩺🧑‍⚕️\"],\n",
    "    [\"What is the normal range for blood sugar?\", \"🍬🔢📏\"],\n",
    "    [\"What is a healthy diet?\", \"🥗🍎🍗\"],\n",
    "    [\"How much water should I drink daily?\", \"💧🥤🕒\"],\n",
    "    [\"What are the symptoms of diabetes?\", \"🍬📉😓\"],\n",
    "    [\"How can I improve my sleep quality?\", \"🛏️💤🌙\"],\n",
    "    [\"What exercises are good for the heart?\", \"🏋️‍♂️💓🏃‍♂️\"],\n",
    "    [\"How do I prevent infections?\", \"🧼🤲🦠\"],\n",
    "    [\"What are the benefits of regular exercise?\", \"💪🏃‍♂️😊\"],\n",
    "    [\"How do I manage stress?\", \"😌🧘‍♂️🛀\"],\n",
    "    [\"What are the risks of smoking?\", \"🚬⚠️🫁\"],\n",
    "    [\"How can I boost my immune system?\", \"🛡️🧄🍊\"],\n",
    "    [\"What are the symptoms of high blood pressure?\", \"🔝💉📈\"],\n",
    "    [\"How can I maintain a healthy weight?\", \"⚖️🍏🏃‍♂️\"],\n",
    "    [\"What are the side effects of antibiotics?\", \"💊🤢😴\"],\n",
    "    [\"How do I know if I have an allergy?\", \"🌸😷🤧\"],\n",
    "    [\"What is a balanced diet?\", \"🍎🥗🍗\"],\n",
    "    [\"How can I lower my cholesterol?\", \"🧈📉🍵\"],\n",
    "    [\"What is mental health?\", \"🧠❤️😊\"],\n",
    "    [\"How do I treat a burn?\", \"🔥🚑🧴\"],\n",
    "    [\"What are the causes of insomnia?\", \"🌙🚫💤\"],\n",
    "    [\"How can I quit smoking?\", \"🚬❌👍\"],\n",
    "    [\"What are the symptoms of depression?\", \"😞🔵😭\"],\n",
    "    [\"How do I handle a panic attack?\", \"😱🧘‍♂️🛏️\"],\n",
    "    [\"What are the risks of obesity?\", \"⚖️⚠️❤️\"],\n",
    "    [\"How can I stay fit at home?\", \"🏠🏋️‍♂️🧘‍♂️\"],\n",
    "    [\"What is a healthy BMI?\", \"⚖️🔢🧍‍♂️\"],\n",
    "    [\"How do I deal with anxiety?\", \"😟🧘‍♂️💤\"],\n",
    "    [\"What are the benefits of yoga?\", \"🧘‍♂️😊💪\"],\n",
    "    [\"How do I treat a sprained ankle?\", \"🦶🧊🛏️\"],\n",
    "    [\"What are the symptoms of a stroke?\", \"🧠🚨😵\"],\n",
    "    [\"How can I manage chronic pain?\", \"💊🔄🧘‍♂️\"],\n",
    "    [\"What is the best way to stay hydrated?\", \"💧🥤🕒\"],\n",
    "    [\"How do I strengthen my bones?\", \"🦴💪🥛\"],\n",
    "    [\"What are the effects of dehydration?\", \"🥵💧🚫\"],\n",
    "    [\"How can I improve my posture?\", \"🧍‍♂️⬆️🧘‍♂️\"],\n",
    "    [\"What is the importance of fiber in diet?\", \"🌾🍎🥗\"],\n",
    "    [\"How do I manage arthritis pain?\", \"🦴🤕💊\"],\n",
    "    [\"What are the symptoms of COVID-19?\", \"🤒🤧😷\"],\n",
    "    [\"How can I prevent back pain?\", \"🪑⬆️🏋️‍♂️\"],\n",
    "    [\"What is the importance of vaccines?\", \"💉🛡️🦠\"],\n",
    "    [\"How do I handle a nosebleed?\", \"👃🩸🧴\"],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d712e891c09e447d913574dd4bcc9ee0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a586c3f3c2204a938fb4312101269159",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/1.29k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d48bb5a670240a3b79fb74fcdea7d99",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.model:   0%|          | 0.00/500k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "700e25bbd24f405786afe7caa9707b14",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/551 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b535fe10f44b44cc9c20183540729c88",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.84M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompt_no_input_template = \"\"\"\\n<|user|>:%s\\n<|assistant|>:\"\"\"\n",
    "\n",
    "model_name_or_path = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n",
    "model = transformers.AutoModelForCausalLM.from_pretrained(\n",
    "    model_name_or_path, torch_dtype=torch.bfloat16, device_map=\"cpu\")\n",
    "\n",
    "# get tokenizer\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained(\n",
    "    model_name_or_path, model_max_length=2048,\n",
    "    padding_side=\"right\", use_fast=False)\n",
    "tokenizer.pad_token = tokenizer.unk_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable intervention params: 16,388 || trainable model params: 0\n",
      "model params: 1,100,048,384 || trainable%: 0.001489752654370519\n"
     ]
    }
   ],
   "source": [
    "# get reft model\n",
    "reft_config = pyreft.ReftConfig(representations={\n",
    "    \"layer\": 8, \"component\": \"block_output\",\n",
    "    \"low_rank_dimension\": 4,\n",
    "    \"intervention\": pyreft.LoreftIntervention(embed_dim=model.config.hidden_size,\n",
    "    low_rank_dimension=4)})\n",
    "reft_model = pyreft.get_reft_model(model, reft_config)\n",
    "reft_model.set_device(\"cpu\")\n",
    "reft_model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_module = pyreft.make_last_position_supervised_data_module(\n",
    "    tokenizer, model, [prompt_no_input_template % e[0] for e in training_examples],\n",
    "    [e[1] for e in training_examples])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa9c2de405274c9ba708cad6999eb71c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "training_args = transformers.TrainingArguments(\n",
    "    num_train_epochs=100.0, output_dir=\"./tmp\", per_device_train_batch_size=10,\n",
    "    learning_rate=4e-3, logging_steps=40, report_to=[])\n",
    "trainer = pyreft.ReftTrainerForCausalLM(\n",
    "    model=reft_model, tokenizer=tokenizer, args=training_args, **data_module)\n",
    "_ = trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}