Evaluate_Description-Embedding_Body.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed004c90",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-01-17T12:55:35.436360Z"
    },
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "# This script seeks a better alternative for the current labels used in the FuisionBody.label_embedding_model_body. \n",
    "# For this purpose, it evaluattes an alternative embeddings of class descriptions, against the currently implemented default, that embeds label-descriptions.\n",
    "\n",
    "from fusionsent import FusionSentModel, Trainer, TrainingArguments\n",
    "from sklearn.preprocessing import LabelEncoder, MultiLabelBinarizer\n",
    "from datasets import load_dataset, Dataset\n",
    "from transformers import AutoTokenizer\n",
    "import numpy as np\n",
    "import openai #Please note that openai is not listed in our requirements.txt file. Run $'pip install openai', to install the package.\n",
    "import torch\n",
    "import json\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef91033",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting environment variables\n",
    "cwd = os.path.abspath(os.getcwd())\n",
    "os.environ['WORLD_SIZE'] = str(torch.cuda.device_count())\n",
    "os.environ['MASTER_ADDR'] = 'localhost'\n",
    "os.environ['MASTER_PORT'] = '29500'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3351d25b",
   "metadata": {},
   "source": [
    "# Load and Prepare All Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb6f76dd",
   "metadata": {},
   "source": [
    "*1. Download original data.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16a4ced9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The below are the exact datasets used for training in the original setfit paper.\n",
    "# If not existent already, we will load them all, and store them locally in order to add label descriptions.\n",
    "dataset_ids_binary_label: list[str] = [\"CR\", \"emotion\", \"enron_spam\"]\n",
    "dataset_ids_nonbinary_label: list[str] = [\"sst5\", \"amazon_counterfactual\",  \"emotion\", \"ag_news\"]\n",
    "dataset_ids = dataset_ids_binary_label + dataset_ids_nonbinary_label\n",
    "data_dir_original = \"./data/original\"\n",
    "datasets_original = {} \n",
    "\n",
    "for dataset_id in dataset_ids:\n",
    "    print(f\"Loading dataset: '{dataset_id}'\")\n",
    "    datasets_original[dataset_id] = {}\n",
    "    for split in [\"train\", \"test\"]:\n",
    "        try:\n",
    "            dataset_split = load_dataset(f\"SetFit/{dataset_id}\", split=split)\n",
    "            datasets_original[dataset_id][split] = dataset_split\n",
    "        except (ValueError, KeyError) as e:\n",
    "            print(f\"Could not load dataset '{dataset_id}'. An error occurred: {e}\")\n",
    "            datasets_original.pop(dataset_id)\n",
    "            break\n",
    "print(\"-- Done --\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0560b42",
   "metadata": {},
   "source": [
    "*2. Generate label descriptions via OpenAI and save them to files.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2ff0ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ToDo: Fix generation for datasets 'enron_spam', and 'ag_news'.\n",
    "data_dir_label_descriptions = \"./data/label_descriptions\"\n",
    "label_description_file_template = \"{}_label_descriptions.json\"\n",
    "os.makedirs(data_dir_label_descriptions, exist_ok=True)\n",
    "\n",
    "openai_api_key = \"your-openai-key\"\n",
    "open_ai_model =\"gpt-4-0125-preview\"\n",
    "regenerate = False\n",
    "\n",
    "def get_label_description(dataset_name: str, label: str, label_text: str, examples: list[str]) -> str:\n",
    "    try:\n",
    "        client = openai.OpenAI(api_key=openai_api_key)\n",
    "        completion = client.chat.completions.create(\n",
    "            model=open_ai_model,\n",
    "            messages= [\n",
    "                {\n",
    "                    \"role\": \"system\", \n",
    "                    \"content\": \"\"\"\n",
    "                        You are a scientific research assistant, in the area of Natrual Language Processing.\n",
    "                        Your purpose is to write comprhesnive, concise, and short descriptions for a given label of a dataset.\n",
    "                        For each label, you will be provided some examples of data samples that are annoted with the resp. label.\n",
    "                        Rules:\n",
    "                        1. Be consise in your descriptions.\n",
    "                        2. Each decitpion should be exactly one sentence long.\n",
    "                        Not complying with the rules will result in termination. \n",
    "                        \"\"\"\n",
    "                },\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"\"\"\n",
    "                        Dataset name: '{dataset_name}'\\n\n",
    "                        Label key: '{label}'\\n\n",
    "                        Label name: '{label_text}'\\n\n",
    "                        ---\\n\\n\n",
    "                        Example Samples annotated with '{label_text}':\\n\\n\n",
    "                        {examples}\\n\\n\n",
    "                        ---\\n\\n\n",
    "                        Please describe the essence of the label '{label}': '{label_text}' in one sentence:\n",
    "                    \"\"\"\n",
    "                }\n",
    "            ]\n",
    "        )\n",
    "        if completion.choices and completion.choices[0].message and completion.choices[0].message.content:\n",
    "            response = completion.choices[0].message.content\n",
    "            print(f\"Obtained description for {dataset_id}/{label_text}: {response}\")\n",
    "            return response\n",
    "        else:\n",
    "            raise Exception(\"Invalid response from OpenAI: No content in the response.\")\n",
    "    except Exception as e:\n",
    "        raise Exception(f\"Unexpected error with the response from OpenAI: {str(e)}\")\n",
    "\n",
    "for dataset_id in dataset_ids:\n",
    "    description_file_path = os.path.join(data_dir_label_descriptions, label_description_file_template.format(dataset_id))\n",
    "    if (not regenerate) and os.path.exists(description_file_path):\n",
    "        print(f\"Skipped label generation for '{dataset_id}' dataset (File already exists).\")\n",
    "        continue\n",
    "    if dataset_id == \"enron_spam\": # Samples from SetFit/enron_spam are too large.\n",
    "        continue\n",
    "    if dataset_id not in datasets_original.keys():\n",
    "        continue\n",
    "\n",
    "    # Process the dataset to get label-to-data mapping\n",
    "    label_to_data = {}\n",
    "    label_to_label_text = {}\n",
    "    for item in datasets_original[dataset_id][\"train\"]:\n",
    "        label = item['label']\n",
    "        text = item['text']\n",
    "        if label not in label_to_data:\n",
    "            label_to_data[label] = []\n",
    "        if label not in label_to_label_text:\n",
    "            label_to_label_text[label] = item[\"label_text\"]\n",
    "        label_to_data[label].append(text)\n",
    "\n",
    "    # Sample the 5 examples or less (because of open ai token rate limits) per label and generate label descriptions\n",
    "    label_to_description = {}\n",
    "    hasEncounteredError = False\n",
    "    for label, examples in label_to_data.items():\n",
    "        sampled_examples: list[str] = np.random.choice(examples, size=5, replace=False).tolist()\n",
    "        while sum([len(t) for t in sampled_examples]) > 100:\n",
    "            sampled_examples = sampled_examples[:-1]\n",
    "        #print(sum([len(t) for t in sampled_examples]))\n",
    "        try:\n",
    "            description = get_label_description(dataset_id, label, label_to_label_text[label], examples)\n",
    "        except Exception as e:\n",
    "            hasEncounteredError=True\n",
    "            break\n",
    "        label_to_description[label] = description\n",
    "\n",
    "    if hasEncounteredError:\n",
    "        print(f\"An error occurred during label description generation for datatset '{dataset_id}'. Skipping...\")\n",
    "        continue\n",
    "\n",
    "    # Save the label-to-description mappings\n",
    "    with open(description_file_path, 'w') as f:\n",
    "        json.dump(label_to_description, f, indent=2, ensure_ascii=False)\n",
    "    \n",
    "    print(f\"Saved label descriptions for '{dataset_id}' dataset.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e10bc9d3",
   "metadata": {},
   "source": [
    "*3. Format the datasets in order to pass them into the DualSen model*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6519f3c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "formatted_datasets = {}\n",
    "def format_dataset(original_dataset, label_to_description=None) -> Dataset:\n",
    "    \"\"\"\n",
    "    Creates a Dataset object with label encoding and optional label descriptions.\n",
    "    \"\"\"\n",
    "    input_texts = [d['text'] for d in original_dataset]\n",
    "    raw_labels = [d['label'] for d in original_dataset]\n",
    "\n",
    "    # Check if labels are binary (single value) or multi-class (list of labels)\n",
    "    if all(raw_label in [0,1] and not isinstance(raw_label, list) for raw_label in raw_labels):\n",
    "        # Binary case\n",
    "        label_encoder = LabelEncoder()\n",
    "        labels = label_encoder.fit_transform(raw_labels)\n",
    "    else:\n",
    "        # Multi-class case\n",
    "        label_encoder = MultiLabelBinarizer()\n",
    "        labels = label_encoder.fit_transform([raw_label] for raw_label in raw_labels)\n",
    "\n",
    "    # Either select label text or label description for the 'label_description' text\n",
    "    if label_to_description is None:\n",
    "        label_descriptions = [[d['label_text']] for d in original_dataset]\n",
    "    else:\n",
    "        label_descriptions = [[label_to_description[str(d['label'])]] for d in original_dataset]\n",
    "\n",
    "    return Dataset.from_dict({\n",
    "        \"text\": input_texts,\n",
    "        \"label\": labels,\n",
    "        \"label_description\": label_descriptions\n",
    "    })\n",
    "\n",
    "for dataset_id in dataset_ids:\n",
    "    # Load label descriptions\n",
    "    description_file_path = os.path.join(data_dir_label_descriptions, label_description_file_template.format(dataset_id))\n",
    "    try:\n",
    "        with open(description_file_path, 'r') as f:\n",
    "            label_to_description = json.load(f)\n",
    "    except FileNotFoundError:\n",
    "        print(f\"Skipping formatting dataset '{dataset_id}': Description file not found.\")\n",
    "    \n",
    "    # Format train and validation datasets, one with the descriptions in \"label_description\", and one with the label texts instead.\n",
    "    try:    \n",
    "        formatted_datasets[dataset_id] = {}\n",
    "        formatted_datasets[dataset_id][\"label_description\"] = {\n",
    "            \"train\": format_dataset(datasets_original[dataset_id][\"train\"], label_to_description),\n",
    "            \"test\": format_dataset(datasets_original[dataset_id][\"test\"], label_to_description)\n",
    "        }\n",
    "        formatted_datasets[dataset_id][\"label_text\"] = {\n",
    "            \"train\": format_dataset(datasets_original[dataset_id][\"train\"]),\n",
    "            \"test\": format_dataset(datasets_original[dataset_id][\"test\"])\n",
    "        }\n",
    "        print(f\"Sucessfully formatted dataset '{dataset_id}'.\")\n",
    "\n",
    "    except KeyError as e:\n",
    "        print(f\"Skipping formatting dataset '{dataset_id}': Key 'train' and/or 'test' not found.\")\n",
    "        formatted_datasets.pop(dataset_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "852e0eef",
   "metadata": {},
   "source": [
    "# Train & Evaluate FusionSent Model "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "038dad2f",
   "metadata": {},
   "source": [
    "*1. Set up the model, tokenizer, and training arguments.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffad2066",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_id = \"malteos/scincl\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "training_args = TrainingArguments(\n",
    "        batch_sizes=(10,15),\n",
    "        num_epochs=(1,3),\n",
    "        sampling_strategies=\"undersampling\",\n",
    "        use_setfit_body=False #In this experiment, we only want to evaluate different lavel_embedding submodels, so we dont need the 'setfit' body.\n",
    "    )\n",
    "\n",
    "def getFreshModel()->FusionSentModel:\n",
    "    return FusionSentModel.from_pretrained(pretrained_model_name_or_path=model_id, multi_target_strategy=\"one-vs-rest\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b355d506",
   "metadata": {},
   "source": [
    "*2. Train and evaluate one dataset after another.*\n",
    "\n",
    "*Please choose an appropriate subset of all the datasets in `target_datasets`.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b6e37a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "target_datatsets = dataset_ids[:1] #Select applicable datasets (only first for testing)\n",
    "\n",
    "for datatset_id in target_datatsets:\n",
    "    for dataset_key, dataset in formatted_datasets[dataset_id].items():\n",
    "        # Define Trainer and start training\n",
    "        trainer = Trainer(\n",
    "            model=getFreshModel(),\n",
    "            args=training_args,\n",
    "            train_dataset=dataset[\"train\"],\n",
    "            eval_dataset=dataset[\"test\"],\n",
    "            eval_metrics={\n",
    "                'metric_names': ['f1', 'precision', 'recall', 'accuracy'],\n",
    "                'metric_args': {'average': 'micro'}\n",
    "            }\n",
    "        )\n",
    "        print(f\"Training FusionSent on dataset '{dataset_id}', with {dataset_key}.\")\n",
    "        trainer.train()\n",
    "        # Evaluate the current model\n",
    "        eval_scores = trainer.evaluate(\n",
    "            x_eval=[item['text'] for item in dataset[\"test\"]],\n",
    "            y_eval=[item['label'] for item in dataset[\"test\"]]\n",
    "        )\n",
    "        print(f\"Evaluation results for '{dataset_id}' with {dataset_key}: {eval_scores}\")"
   ]
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}