create_evaluation_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset, load_dataset, concatenate_datasets\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HotpotQA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hotpotqa_dataset = load_dataset(\"hotpot_qa\", \"distractor\")\n",
    "\n",
    "# let's sample a few examples from each level (of difficulty) and type (comparion or bridge)\n",
    "dataset_df = pd.DataFrame(hotpotqa_dataset[\"train\"])\n",
    "sample_indicies = (\n",
    "    dataset_df.groupby([\"level\", \"type\"]).sample(4, random_state=10).index.values\n",
    ")\n",
    "hotpotqa_dataset.reset_format()\n",
    "hotpotqa_dataset_leftout = hotpotqa_dataset[\"train\"].select(\n",
    "    [i for i in range(len(hotpotqa_dataset[\"train\"])) if i not in sample_indicies]\n",
    ")\n",
    "hotpotqa_dataset = hotpotqa_dataset[\"train\"].select(sample_indicies)\n",
    "\n",
    "hotpotqa_dataset_leftout_df = pd.DataFrame(hotpotqa_dataset_leftout)\n",
    "hotpotqa_dataset_leftout_df = (\n",
    "    hotpotqa_dataset_leftout_df.groupby([\"level\", \"type\"])\n",
    "    .sample(6, random_state=42)\n",
    "    .reset_index(drop=True)\n",
    ")\n",
    "hotpotqa_dataset_leftout = Dataset.from_pandas(hotpotqa_dataset_leftout_df)\n",
    "\n",
    "task_column = [f\"HotpotQA-{level}\" for level in hotpotqa_dataset[\"level\"]]\n",
    "hotpotqa_dataset = hotpotqa_dataset.add_column(\"task\", task_column).select_columns(\n",
    "    [\"question\", \"answer\", \"task\"]\n",
    ")\n",
    "\n",
    "task_column = [f\"HotpotQA-{level}\" for level in hotpotqa_dataset_leftout[\"level\"]]\n",
    "hotpotqa_dataset_leftout = hotpotqa_dataset_leftout.add_column(\n",
    "    \"task\", task_column\n",
    ").select_columns([\"question\", \"answer\", \"task\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hotpotqa_dataset = concatenate_datasets([hotpotqa_dataset, hotpotqa_dataset_leftout])\n",
    "print(len(hotpotqa_dataset), len(pd.Series(hotpotqa_dataset[\"question\"]).unique()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GSM8K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "\n",
    "math_dataset = load_dataset(\"gsm8k\", \"main\")[\"train\"]\n",
    "\n",
    "first_selection = np.random.randint(0, len(math_dataset), 15)\n",
    "second_selection = np.random.randint(0, len(math_dataset), 15)\n",
    "second_selection_first_excluded = [\n",
    "    i for i in second_selection if i not in first_selection\n",
    "][:20]\n",
    "\n",
    "math_dataset = math_dataset.select(\n",
    "    list(first_selection) + list(second_selection_first_excluded)[:5]\n",
    ")\n",
    "\n",
    "\n",
    "task_column = [\"GSM8K\"] * len(math_dataset)\n",
    "math_dataset = math_dataset.add_column(\"task\", task_column).select_columns(\n",
    "    [\"question\", \"answer\", \"task\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GAIA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gaia_dataset = load_dataset(\"gaia-benchmark/GAIA\", \"2023_level1\")[\"validation\"]\n",
    "gaia_dataset.set_format(\"pandas\")\n",
    "gaia_dataset_df = gaia_dataset[:]\n",
    "gaia_dataset_df[\"number_of_steps\"] = gaia_dataset_df[\"Annotator Metadata\"].apply(\n",
    "    lambda row: int(row[\"Number of steps\"])\n",
    ")\n",
    "gaia_dataset_df[\"tools_used\"] = gaia_dataset_df[\"Annotator Metadata\"].apply(\n",
    "    lambda row: row[\"Tools\"]\n",
    ")\n",
    "gaia_dataset_df = gaia_dataset_df.loc[\n",
    "    ~gaia_dataset_df[\"tools_used\"]\n",
    "    .str.lower()\n",
    "    .str.contains(\n",
    "        \"pdf|excel|image|video|parsing|audio|word|file|speech|viewer|markdown|python|editor|model\"\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.set_option(\"display.max_colwidth\", None)\n",
    "gaia_dataset_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_indicies = [\n",
    "    0,\n",
    "    2,\n",
    "    11,\n",
    "    12,\n",
    "    23,\n",
    "    28,\n",
    "    25,\n",
    "    29,\n",
    "    32,\n",
    "    35,\n",
    "    36,\n",
    "    37,\n",
    "    39,\n",
    "    40,\n",
    "    41,\n",
    "    42,\n",
    "    43,\n",
    "    47,\n",
    "    49,\n",
    "    52,\n",
    "]\n",
    "print(len(selected_indicies))\n",
    "gaia_dataset = gaia_dataset.rename_columns(\n",
    "    {\"Question\": \"question\", \"Final answer\": \"answer\"}\n",
    ").select_columns([\"question\", \"answer\"])\n",
    "gaia_dataset.reset_format()\n",
    "gaia_dataset = gaia_dataset.select(selected_indicies)\n",
    "\n",
    "task_column = [\"GAIA\"] * len(gaia_dataset)\n",
    "gaia_dataset = gaia_dataset.add_column(\"task\", task_column)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_eval_dataset = concatenate_datasets([math_dataset, hotpotqa_dataset, gaia_dataset])\n",
    "pd.Series(full_eval_dataset[\"task\"]).value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_eval_dataset.push_to_hub(\"m-ric/agents_small_benchmark\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml2",
   "language": "python",
   "name": "ml2"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}