{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ***Setup***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from ollama import Client\n",
    "\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "parent_dir = os.path.abspath('..')\n",
    "sys.path.append(parent_dir)\n",
    "from constants import CONSTANTS\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ***Read dataset***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(CONSTANTS.DATASET_PATH)\n",
    "assert df.shape == (1200, 3), f\"Expected shape (1200, 3), but got {df.shape}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ***Some helper functions***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ***Send prompts to `llama3` API***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(host='http://172.25.1.141:11434')\n",
    "\n",
    "def clean_generated_text(text):\n",
    "    cleaned_text = re.sub(r'[^\\w\\s,.\\']', '', text)  # Remove special characters\n",
    "    cleaned_text = cleaned_text.replace('\\n', ' ')  # replace newlines with spaces\n",
    "    cleaned_text = ' '.join(cleaned_text.split())  # eemove extra spaces\n",
    "    return cleaned_text.strip()  # Trim leading and trailing spaces\n",
    "\n",
    "\n",
    "def send_prompt(prompt):\n",
    "    response = client.chat(model='llama3', messages=[\n",
    "        {\n",
    "            'role': 'user',\n",
    "            'content': prompt + '\\nSend the response only.',\n",
    "        },\n",
    "    ])\n",
    "    \n",
    "    # Extracting the content from the response\n",
    "    return clean_generated_text(response['message']['content'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ***Example usage***\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I've noticed a persistent skin rash on my arms, legs, and torso over the past few weeks. The affected areas appear red, itchy, and are characterized by dry, scaling patches.\n",
      "For the past few weeks, I've been dealing with a persistent skin rash that affects my arms, legs, and torso. The rash appears as red, itchy areas with dry, flaky patches.\n",
      "For several weeks, I've been dealing with a persistent skin issue affecting my arms, legs, and torso. The rash appears as red, itchy areas covered in dry, flaky scales.\n",
      "For several weeks, I've noticed a persistent skin issue affecting my arms, legs, and torso. The affected areas appear red, irritated, and feature dry, flaky scales.\n",
      "For the past few weeks, I've been dealing with an uncomfortable skin issue that's affected my arms, legs, and torso. The rash has characterized by a reddish hue, intense itchiness, and the presence of dry, flaky scales.\n"
     ]
    }
   ],
   "source": [
    "text = \"I have been experiencing a skin rash on my arms, legs, and torso for the past few weeks. It is red, itchy, and covered in dry, scaly patches.\"\n",
    "for i in range(5):\n",
    "  result_i = send_prompt(f'Rephrase the following text:\\n{text}')\n",
    "  print(result_i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ***Augment dataset using `llm`***\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**`n_augmentations` is number of rephrased texts to generate for each row in the original dataset**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_dataset(df, text_column=\"text\", label_column=\"label\", n_augmentations=5):\n",
    "    augmented_texts = []\n",
    "    augmented_labels = []\n",
    "    total_rows = len(df)\n",
    "    \n",
    "    \n",
    "    if total_rows < 10:\n",
    "        progress_step = 1\n",
    "    else:\n",
    "        progress_step = total_rows // 10\n",
    "\n",
    "    for index, row in df.iterrows():\n",
    "        text, label = row[text_column], row[label_column]\n",
    "        for _ in range(n_augmentations):\n",
    "            augmented_text = send_prompt(f'Rephrase the following text:\\n{text}')\n",
    "            augmented_texts.append(augmented_text)\n",
    "            augmented_labels.append(label)\n",
    "        \n",
    "        # Verbose section\n",
    "        if (index + 1) % progress_step == 0:\n",
    "            percentage_done = (index + 1) / total_rows * 100\n",
    "            print(f\"Done {percentage_done:.0f}%\")\n",
    "\n",
    "    augmented_df = pd.DataFrame({text_column: augmented_texts, label_column: augmented_labels})\n",
    "    return pd.concat([df, augmented_df], ignore_index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ***Augment dataset***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Done 10%\n",
      "Done 20%\n",
      "Done 30%\n",
      "Done 40%\n",
      "Done 50%\n",
      "Done 60%\n",
      "Done 70%\n",
      "Done 80%\n",
      "Done 90%\n",
      "Done 100%\n"
     ]
    }
   ],
   "source": [
    "augmented_df = augment_dataset(df, n_augmentations=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "augmented_df = augmented_df.drop(columns=['Unnamed: 0'])\n",
    "augmented_df['Id'] = range(len(augmented_df))\n",
    "augmented_df = augmented_df[['Id', 'label', 'text']]\n",
    "augmented_df.to_csv(CONSTANTS.AUGMENTED_DATASET_PATH, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ***Preprocess augmented dataset***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        Id      label                                               text\n",
      "0        0  Psoriasis  I have been experiencing a skin rash on my arm...\n",
      "1        1  Psoriasis  My skin has been peeling, especially on my kne...\n",
      "2        2  Psoriasis  I have been experiencing joint pain in my fing...\n",
      "3        3  Psoriasis  There is a silver like dusting on my skin, esp...\n",
      "4        4  Psoriasis  My nails have small dents or pits in them, and...\n",
      "...    ...        ...                                                ...\n",
      "1195  1195   diabetes  I'm shaking and trembling all over. I've lost ...\n",
      "1196  1196   diabetes  Particularly in the crevices of my skin, I hav...\n",
      "1197  1197   diabetes  I regularly experience these intense urges and...\n",
      "1198  1198   diabetes  I have trouble breathing, especially outside. ...\n",
      "1199  1199   diabetes  I constantly sneeze and have a dry cough. My i...\n",
      "\n",
      "[1200 rows x 3 columns]\n",
      "      Unnamed: 0      label                                               text\n",
      "0              0  Psoriasis  I have been experiencing a skin rash on my arm...\n",
      "1              1  Psoriasis  My skin has been peeling, especially on my kne...\n",
      "2              2  Psoriasis  I have been experiencing joint pain in my fing...\n",
      "3              3  Psoriasis  There is a silver like dusting on my skin, esp...\n",
      "4              4  Psoriasis  My nails have small dents or pits in them, and...\n",
      "...          ...        ...                                                ...\n",
      "1195         295   diabetes  I'm shaking and trembling all over. I've lost ...\n",
      "1196         296   diabetes  Particularly in the crevices of my skin, I hav...\n",
      "1197         297   diabetes  I regularly experience these intense urges and...\n",
      "1198         298   diabetes  I have trouble breathing, especially outside. ...\n",
      "1199         299   diabetes  I constantly sneeze and have a dry cough. My i...\n",
      "\n",
      "[1200 rows x 3 columns]\n"
     ]
    }
   ],
   "source": [
    "from sklearn.utils import shuffle\n",
    "from sklearn.model_selection import train_test_split\n",
    "original_data = augmented_df.head(1200)\n",
    "print(original_data)\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ***1. Split ORIGINAL DATASET into train and test sets***\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train set shape: (720, 3)\n",
      "Test set shaape: (480, 3)\n"
     ]
    }
   ],
   "source": [
    "train_df, test_df = train_test_split(original_data, test_size=0.4, random_state=42)\n",
    "print(f\"Train set shape: {train_df.shape}\")\n",
    "print(f\"Test set shaape: {test_df.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ***2. Add augmented texts to TRAIN SET only***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_augmented_texts_ids(original_id):\n",
    "    start = 1200\n",
    "    shift = 5 * original_id\n",
    "    augmented_texts_ids = []\n",
    "    for i in range(5):\n",
    "        augmented_texts_ids.append(start + shift + i)\n",
    "    return augmented_texts_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New train_set shape: (4320, 3)\n"
     ]
    }
   ],
   "source": [
    "original_ids = train_df[\"Id\"]\n",
    "dataframes_to_concat = []\n",
    "\n",
    "for id in original_ids:\n",
    "    for augmented_id in get_augmented_texts_ids(id):\n",
    "        augmented_rows = augmented_df[augmented_df[\"Id\"] == augmented_id]\n",
    "        dataframes_to_concat.append(augmented_rows)\n",
    "\n",
    "augmented_data = pd.concat(dataframes_to_concat, ignore_index=True)\n",
    "train_df = pd.concat([train_df, augmented_data], ignore_index=True)\n",
    "\n",
    "print(f\"New train_set shape: {train_df.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ***3- Count unique texts for each label***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Count of unique text values for each label in train set\n",
      "label\n",
      "Fungal infection                   144\n",
      "Varicose Veins                     150\n",
      "Migraine                           155\n",
      "urinary tract infection            162\n",
      "Impetigo                           162\n",
      "Jaundice                           164\n",
      "Dengue                             166\n",
      "Pneumonia                          168\n",
      "Acne                               172\n",
      "Cervical spondylosis               173\n",
      "Common Cold                        174\n",
      "gastroesophageal reflux disease    174\n",
      "Arthritis                          179\n",
      "Dimorphic Hemorrhoids              184\n",
      "Bronchial Asthma                   185\n",
      "Hypertension                       191\n",
      "peptic ulcer disease               192\n",
      "Typhoid                            192\n",
      "Malaria                            194\n",
      "diabetes                           197\n",
      "Psoriasis                          198\n",
      "drug reaction                      204\n",
      "Chicken pox                        210\n",
      "allergy                            210\n",
      "Name: text, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Train set\n",
    "unique_text_counts = train_df.groupby('label')['text'].nunique()\n",
    "unique_text_counts = unique_text_counts.sort_values()\n",
    "\n",
    "print(\"Count of unique text values for each label in train set\")\n",
    "print(unique_text_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Count of unique text values for each label in test set\n",
      "label\n",
      "Chicken pox                        15\n",
      "allergy                            15\n",
      "drug reaction                      16\n",
      "diabetes                           17\n",
      "Dimorphic Hemorrhoids              17\n",
      "Malaria                            17\n",
      "Psoriasis                          17\n",
      "Hypertension                       18\n",
      "Typhoid                            18\n",
      "peptic ulcer disease               18\n",
      "Arthritis                          19\n",
      "Bronchial Asthma                   19\n",
      "Jaundice                           20\n",
      "gastroesophageal reflux disease    21\n",
      "Acne                               21\n",
      "Common Cold                        21\n",
      "Cervical spondylosis               21\n",
      "Pneumonia                          21\n",
      "Dengue                             22\n",
      "Migraine                           23\n",
      "Impetigo                           23\n",
      "urinary tract infection            23\n",
      "Varicose Veins                     25\n",
      "Fungal infection                   26\n",
      "Name: text, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Test set\n",
    "unique_text_counts = test_df.groupby('label')['text'].nunique()\n",
    "unique_text_counts = unique_text_counts.sort_values()\n",
    "\n",
    "print(\"Count of unique text values for each label in test set\")\n",
    "print(unique_text_counts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ***4- Save in .csv files***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.to_csv(CONSTANTS.AUGMENTED_TRAIN_SET_PATH, index=False)\n",
    "test_df.to_csv(CONSTANTS.AUGMENTED_TEST_SET_PATH, index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "NLP",
   "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.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
