notebooks/upload_human_annotation_csv_to_db.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Uploading rewards from human annotation to DB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Import neccessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from sotopia.database import EpisodeLog\n",
    "from sotopia.envs.evaluators import EvaluationBySocialDimensions\n",
    "from sotopia.database.logs import AnnotationForEpisode\n",
    "import tqdm\n",
    "from pathlib import Path\n",
    "from pandas import read_csv\n",
    "from redis_om import Migrator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find all csv files in the directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csv_files = list(Path(\".\").glob(\"*.csv\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Iterate through all csv files and upload them to DB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SOCIAL_DIMENSIONS: list[str] = list(EvaluationBySocialDimensions.__fields__.keys())\n",
    "\n",
    "Migrator().run()\n",
    "\n",
    "for csv_file in csv_files:\n",
    "    print(f\"Processing {csv_file}\")\n",
    "    d = read_csv(csv_file)\n",
    "    for _, row in tqdm.tqdm(d.iterrows()):\n",
    "        rewards: list[tuple[float, dict[str, float]]] = []\n",
    "        for agent in [\"agent1\", \"agent2\"]:\n",
    "            rewards_for_agent = {\n",
    "                social_dim: row[f\"Answer.{agent}_{social_dim}\"]\n",
    "                for social_dim in SOCIAL_DIMENSIONS\n",
    "            }\n",
    "            rewards.append(\n",
    "                (\n",
    "                    sum(rewards_for_agent.values()) / len(rewards_for_agent),\n",
    "                    rewards_for_agent,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        reasoning = \"\"\n",
    "        for agent_name_in_reasoning, agent in [\n",
    "            (\"agent 1\", \"agent1\"),\n",
    "            (\"agent 2\", \"agent2\"),\n",
    "        ]:\n",
    "            reasoning += f\"{agent_name_in_reasoning} comments: \"\n",
    "            for social_dim in SOCIAL_DIMENSIONS:\n",
    "                reasoning += (\n",
    "                    f\"{social_dim}: {row[f'Answer.{agent}_{social_dim}_rationale']} \"\n",
    "                )\n",
    "\n",
    "        episode_id = row[\"Input.episode_id\"]\n",
    "        assert EpisodeLog.get(pk=episode_id)\n",
    "        existing_annotations = AnnotationForEpisode.find(\n",
    "            (AnnotationForEpisode.episode == episode_id)\n",
    "            & (AnnotationForEpisode.annotator_id == row[\"WorkerId\"])\n",
    "        ).all()\n",
    "        if existing_annotations:\n",
    "            print(\n",
    "                f\"Skipping {episode_id} for {row['WorkerId']} because it already exists\"\n",
    "            )\n",
    "            continue\n",
    "        annotation = AnnotationForEpisode(\n",
    "            episode=episode_id,\n",
    "            annotator_id=row[\"WorkerId\"],\n",
    "            rewards=rewards,\n",
    "            reasoning=reasoning,\n",
    "        )\n",
    "        annotation.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove disqualified workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "worker_id_widget = widgets.Textarea(\n",
    "    placeholder=\"Type something\", description=\"String:\", disabled=False\n",
    ")\n",
    "worker_id_widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layout = widgets.Layout(width=\"auto\", height=\"40px\")\n",
    "worker_id = worker_id_widget.value\n",
    "annotations = AnnotationForEpisode.find(\n",
    "    AnnotationForEpisode.annotator_id == worker_id\n",
    ").all()\n",
    "print(f\"Found {len(annotations)} annotations\")\n",
    "\n",
    "\n",
    "def _f(*args, **kwargs):\n",
    "    for annotation in annotations:\n",
    "        print(\n",
    "            f\"Deleting Worker {worker_id}'s annotation for Episode {annotation.episode}\"\n",
    "        )\n",
    "        annotation.delete(annotation.pk)\n",
    "\n",
    "\n",
    "button = widgets.Button(\n",
    "    description=\"Do you want to delete all annotations for this worker?\",\n",
    "    disabled=False,\n",
    "    button_style=\"warning\",  # 'success', 'info', 'warning', 'danger' or ''\n",
    "    tooltip=\"Yes\",\n",
    "    icon=\"trash\",  # (FontAwesome names without the `fa-` prefix)\n",
    "    layout=layout,\n",
    ")\n",
    "button.on_click(_f)\n",
    "button"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sotopia",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}