argilla_labeller_setfit.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to build a custom text classifier without days of human labeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we will learn how to build a text classifier with AI and human feedback saving time and resources."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting started\n",
    "\n",
    "### Deploy the Argilla server¶\n",
    "\n",
    "If you already have deployed Argilla, you can skip this step. Otherwise, you can quickly deploy Argilla following [this guide](https://docs.argilla.io/latest/getting_started/quickstart/).\n",
    "\n",
    "### Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install argilla\n",
    "!pip install \"distilabel[hf-inference-endpoints]\"\n",
    "!pip install setfit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure the Argilla dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argilla as rg\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Connect to the Argilla server\n",
    "client = rg.Argilla(api_url=\"<http://localhost:6900>\", api_key=\"argilla.apikey\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure the dataset\n",
    "labels = [\"positive\", \"neutral\", \"negative\"]\n",
    "settings = rg.Settings(\n",
    "\tguidelines = \"We are dealing with input data about detailed customer reviews of 3 corresponding labels: positive, negative, and neutral.\",\n",
    "    fields=[\n",
    "        rg.TextField(\n",
    "            name=\"text\",\n",
    "            title=\"Text\",\n",
    "            description=\"Provide a concise response to the prompt\",\n",
    "        )\n",
    "    ],\n",
    "    questions=[\n",
    "        rg.LabelQuestion(\n",
    "            name=\"label\",\n",
    "            title=\"Emotion\",\n",
    "            description=\"Provide a single label for the emotion of the text\",\n",
    "            labels=labels,\n",
    "        )\n",
    "    ],\n",
    "    mapping = {\"labels\": \"label\"}\n",
    ")\n",
    "\n",
    "# Create the dataset\n",
    "dataset_name = \"pc-component\"\n",
    "dataset = rg.Dataset(\n",
    "    name=dataset_name,\n",
    "    settings=settings,\n",
    ").create()\n",
    "\n",
    "# Load the records (in our case, they have suggestions, but it's not required)\n",
    "hf_dataset = load_dataset(\"argilla/pc-components-reviews\", split=\"train\")\n",
    "records = [\n",
    "    rg.Record(\n",
    "        fields = {\"text\": sample[\"text\"]},\n",
    "        suggestions = [rg.Suggestion(\"label\", sample[\"labels\"] if sample[\"labels\"] in labels else \"neutral\")]\n",
    "    ) for sample in hf_dataset\n",
    "]\n",
    "dataset.records.log(records)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Active auto-labeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from collections import Counter, defaultdict\n",
    "from distilabel.steps.tasks import ArgillaLabeller\n",
    "from distilabel.llms.huggingface import InferenceEndpointsLLM\n",
    "\n",
    "\n",
    "example_records_dict = defaultdict(list)\n",
    "counter = Counter()\n",
    "max_samples_per_label = 16\n",
    "\n",
    "# Helper function to get the example records\n",
    "def get_example_records(num_samples_per_label):\n",
    "    example_records = []    \n",
    "    for label, records in example_records_dict.items():\n",
    "        selected_records = records[:num_samples_per_label]\n",
    "        example_records.extend(selected_records)\n",
    "    random.shuffle(example_records)\n",
    "    return example_records\n",
    "\n",
    "# Simulate record annotation for few-shot\n",
    "for record in dataset.records(with_responses=True, with_suggestions=True):\n",
    "    label_value = record.suggestions[\"label\"].value\n",
    "\n",
    "    if counter[label_value] >= max_samples_per_label:\n",
    "        continue\n",
    "    counter[label_value] += 1\n",
    "\n",
    "    record.responses.add(\n",
    "        rg.Response(question_name=\"label\", value=label_value, user_id=client.me)\n",
    "    )\n",
    "    example_records_dict[label_value].append(record.to_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the ArgillaLabeller\n",
    "labeller = ArgillaLabeller(\n",
    "    llm=InferenceEndpointsLLM(\n",
    "        model_id=\"meta-llama/Llama-3.1-8B-Instruct\",\n",
    "        tokenizer_id=\"meta-llama/Llama-3.1-8B-Instruct\",\n",
    "    ),\n",
    "    example_records=get_example_records(1),\n",
    ")\n",
    "labeller.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the loop to start annotating and getting improved suggestions\n",
    "while True:\n",
    "    pending_records = list(\n",
    "        dataset.records(\n",
    "            query=rg.Query(filter=rg.Filter((\"status\", \"==\", \"pending\"))),\n",
    "            limit=1,\n",
    "        )\n",
    "    )\n",
    "    if not pending_records:\n",
    "        sleep(5)\n",
    "        continue\n",
    "\n",
    "    results = next(\n",
    "        labeller.process(\n",
    "            [\n",
    "                {\n",
    "                    \"record\": record,\n",
    "                    \"fields\": dataset.fields,\n",
    "                    \"question\": dataset.questions[0],\n",
    "                    \"guidelines\": dataset.guidelines,\n",
    "                }\n",
    "                for record in pending_records\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    for record, suggestion in zip(pending_records, results):\n",
    "        record.suggestions.add(rg.Suggestion(**suggestion[\"suggestion\"]))\n",
    "\n",
    "    dataset.records.log(pending_records)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train you small classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "from setfit import sample_dataset\n",
    "\n",
    "# Helper function to split the dataset\n",
    "def sample_and_split(dataset, label_column, num_samples):\n",
    "    train_dataset = sample_dataset(\n",
    "        dataset, label_column=label_column, num_samples=num_samples\n",
    "    )\n",
    "    eval_dataset = dataset.filter(lambda x: x[\"id\"] not in set(train_dataset[\"id\"]))\n",
    "    return train_dataset, eval_dataset\n",
    "\n",
    "# Retrieve the data from Argilla\n",
    "annotated_dataset = client.datasets(dataset_name).records.to_datasets()\n",
    "annotated_dataset = Dataset.from_list([\n",
    "    {\"text\": record[\"text\"], \"label\": record[\"label.responses\"][0], \"id\": i} \n",
    "    for i, record in enumerate(annotated_dataset) \n",
    "    if record.get(\"label.responses\") is not None\n",
    "])\n",
    "train_dataset, eval_dataset = sample_and_split(\n",
    "    annotated_dataset, \"label\", 8\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from setfit import SetFitModel, Trainer\n",
    "\n",
    "# Function to train our SetFit model\n",
    "def train_model(model_name, train_dataset, eval_dataset):\n",
    "    model = SetFitModel.from_pretrained(model_name)\n",
    "\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        train_dataset=train_dataset,\n",
    "    )\n",
    "    trainer.train()\n",
    "    results = trainer.evaluate(eval_dataset)\n",
    "\n",
    "    return model, results\n",
    "\n",
    "# Train the classifier\n",
    "model, results = train_model(\n",
    "    model_name=\"TaylorAI/bge-micro-v2\",\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    ")\n",
    "print(results)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "distilabel-exp",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}