notebook/ExampleTokenDatasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Token Datasets\n",
    "\n",
    "Each model is trained on a token sequence dataset stored in a MongoDB instance.\n",
    "These datasets can be accessed by first downloading the archive files listed in the following table and importing them into a MongoDB instance.\n",
    "\n",
    "| File                                                                                                            | Size  | Description                      |\n",
    "| --                                                                                                              | ---   | ---                              |\n",
    "| [maze.gz                   ](https://dl.fbaipublicfiles.com/searchformer/tokenSeqDB/maze.gz                   ) | 25GB  | Maze token training data         |\n",
    "| [maze.vocabulary.gz        ](https://dl.fbaipublicfiles.com/searchformer/tokenSeqDB/maze.vocabulary.gz        ) | 7.2kB | Maze token meta data             |\n",
    "| [sokoban.gz                ](https://dl.fbaipublicfiles.com/searchformer/tokenSeqDB/sokoban.gz                ) | 1.4GB | Sokoban token training data      |\n",
    "| [sokoban.vocabulary.gz     ](https://dl.fbaipublicfiles.com/searchformer/tokenSeqDB/sokoban.vocabulary.gz     ) | 312B  | Sokoban token meta data          |\n",
    "| [searchformer.gz           ](https://dl.fbaipublicfiles.com/searchformer/tokenSeqDB/searchformer.gz           ) | 8.6GB | Searchformer token training data |\n",
    "| [searchformer.vocabulary.gz](https://dl.fbaipublicfiles.com/searchformer/tokenSeqDB/searchformer.vocabulary.gz) | 1.9MB | Searchformer token meta data     |\n",
    "\n",
    "**To properly import and index the dataset, both the token training data itself and the meta data must be imported into MongoDB.**\n",
    "\n",
    "Data can be imported from these files into a live MongoDB instance with [`mongorestore`](https://www.mongodb.com/docs/database-tools/mongorestore/).\n",
    "For example, to import the maze datasets to a MongoDB instance running on localhost with port the default port, run\n",
    "\n",
    "```\n",
    "mongorestore --gzip --archive=maze.vocabulary.gz\n",
    "mongorestore --gzip --archive=maze.gz \n",
    "```\n",
    "\n",
    "Once all imports completed, the imported datasets can be listed by running \n",
    "\n",
    "```\n",
    "python -m searchformer.trace list-token-datasets\n",
    "```\n",
    "\n",
    "\n",
    "The following datasets are included in the files `maze.gz` and `maze.vocabulary.gz`.\n",
    "\n",
    "| Experiment                   | Dataset Name                                  |\n",
    "| ---                          | ---                                           |\n",
    "| 10x10 Maze, deterministic    | `maze.10-by-10-deterministic.simple`          |\n",
    "| 20x20 Maze, deterministic    | `maze.20-by-20-deterministic.simple`          |\n",
    "| 30x30 Maze, deterministic    | `maze.30-by-30-deterministic.simple`          |\n",
    "| 10x10 Maze, nondeterministic | `maze.10-by-10-nondeterministic.simple`       |\n",
    "| 20x20 Maze, nondeterministic | `maze.20-by-20-nondeterministic.simple`       |\n",
    "| 30x30 Maze, nondeterministic | `maze.30-by-30-nondeterministic.simple`       |\n",
    "\n",
    "\n",
    "The following datasets are included in the files `sokoban.gz` and `sokoban.vocabulary.gz`.\n",
    "\n",
    "| Experiment                   | Dataset Name                                  |\n",
    "| ---                          | ---                                           |\n",
    "| Sokoban                      | `sokoban.7-by-7-walls-2-boxes-2.with-box-40k` |\n",
    "\n",
    "\n",
    "The following datasets are included in the files `searchformer.gz` and `searchformer.vocabulary.gz`.\n",
    "\n",
    "| Model        | Step | Repeat | Checkpoint Name                         | Dataset Name                        |\n",
    "| ---          | ---  | ---    | ---                                     | ---                                 |\n",
    "| Searchformer | 1    | 0      | sokoban-7722-m-trace-plan-100k-0-step-1 | `65b8382b9ee4fbaa76e005b7.improved` |\n",
    "| Searchformer | 1    | 1      | sokoban-7722-m-trace-plan-100k-1-step-1 | `65b8398ff7574141c3ba77ae.improved` |\n",
    "| Searchformer | 1    | 2      | sokoban-7722-m-trace-plan-100k-2-step-1 | `65b8495e2382373d6a21ca99.improved` |\n",
    "| Searchformer | 2    | 0      | sokoban-7722-m-trace-plan-100k-0-step-2 | `65ba856d986d307d60c563ca.improved` |\n",
    "| Searchformer | 2    | 1      | sokoban-7722-m-trace-plan-100k-1-step-2 | `65ba8e2678ad82e62025d0c6.improved` |\n",
    "| Searchformer | 2    | 2      | sokoban-7722-m-trace-plan-100k-2-step-2 | `65ba8ee773586ec314e0da96.improved` |\n",
    "| Searchformer | 3    | 0      | sokoban-7722-m-trace-plan-100k-0-step-3 | `65c8b912c3dd164d1eb691a4.improved` |\n",
    "| Searchformer | 3    | 1      | sokoban-7722-m-trace-plan-100k-1-step-3 | `65ca34d2e25a422484c5d3da.improved` |\n",
    "| Searchformer | 3    | 2      | sokoban-7722-m-trace-plan-100k-2-step-3 | `65ca57d67f455f390d05bf33.improved` |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading a token dataset\n",
    "\n",
    "A token dataset is loaded by instantiating an object of the class `searchformer.trace.TokenizedDataset` and passing the dataset name into the class constructor.\n",
    "This object is then used to access the token sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - 2024-04-26 13:56:34 - root - Connecting to mongodb:https://localhost:27017/mongo\n",
      "DEBUG - 2024-04-26 13:56:34 - root - Loading all ids from Collection(Database(MongoClient(host=['localhost:27017'], document_class=dict, tz_aware=False, connect=True, sockettimeoutms=1800000, connecttimeoutms=1800000), 'tokenSeqDB'), 'maze.10-by-10-deterministic.simple.meta.train') ...\n",
      "DEBUG - 2024-04-26 13:56:37 - root - Finished loading.\n",
      "DEBUG - 2024-04-26 13:56:37 - root - Loading all ids from Collection(Database(MongoClient(host=['localhost:27017'], document_class=dict, tz_aware=False, connect=True, sockettimeoutms=1800000, connecttimeoutms=1800000), 'tokenSeqDB'), 'maze.10-by-10-deterministic.simple.meta.test') ...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of train sequences: 1000000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG - 2024-04-26 13:56:38 - root - Finished loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of test sequences:  100000\n"
     ]
    }
   ],
   "source": [
    "import sys; sys.path.append(\"..\")\n",
    "\n",
    "import logging\n",
    "from searchformer.trace import TokenizedDataset\n",
    "\n",
    "\n",
    "logging.basicConfig(\n",
    "    level=logging.DEBUG,\n",
    "    format=\"%(levelname)s - %(asctime)s - %(name)s - %(message)s\",\n",
    "    datefmt=\"%Y-%m-%d %H:%M:%S\",\n",
    ")\n",
    "\n",
    "\n",
    "tok_dataset = TokenizedDataset(\"maze.10-by-10-deterministic.simple\")\n",
    "print(f\"Number of train sequences: {len(tok_dataset.train_ids)}\")\n",
    "print(f\"Number of test sequences:  {len(tok_dataset.test_ids)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A tokenized sequence is represented with the data class `searchformer.trace.TokenizedTrace`.\n",
    "This class is returned by the training or test iterators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG - 2024-04-26 13:56:38 - root - Loading all ids from Collection(Database(MongoClient(host=['localhost:27017'], document_class=dict, tz_aware=False, connect=True, sockettimeoutms=1800000, connecttimeoutms=1800000), 'tokenSeqDB'), 'maze.10-by-10-deterministic.simple.meta.train') ...\n",
      "DEBUG - 2024-04-26 13:56:41 - root - Finished loading.\n",
      "DEBUG - 2024-04-26 13:56:41 - root - Iterating over 1000000 ids.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt:          \n",
      "\tstart 1 3 \n",
      "\tgoal  9 4 \n",
      "\twall  0 0 \n",
      "\twall  3 0 \n",
      "\twall  4 0 \n",
      "\twall  6 0 \n",
      "\twall  7 0 \n",
      "\twall  6 1 \n",
      "\twall  8 1 \n",
      "\twall  1 2 \n",
      "\twall  3 2 \n",
      "\twall  6 2 \n",
      "\twall  7 2 \n",
      "\twall  2 3 \n",
      "\twall  4 3 \n",
      "\twall  5 3 \n",
      "\twall  6 3 \n",
      "\twall  7 3 \n",
      "\twall  7 4 \n",
      "\twall  8 4 \n",
      "\twall  1 5 \n",
      "\twall  2 5 \n",
      "\twall  4 5 \n",
      "\twall  5 6 \n",
      "\twall  6 6 \n",
      "\twall  0 7 \n",
      "\twall  2 7 \n",
      "\twall  5 7 \n",
      "\twall  8 8 \n",
      "\twall  1 9 \n",
      "\twall  7 9 \n",
      "\twall  8 9 \n",
      "\twall  9 9\n",
      "Execution trace: \n",
      "\tcreate 1 3 c0 c9 \n",
      "\tclose  1 3 c0 c9 \n",
      "\tcreate 1 4 c1 c8 \n",
      "\tcreate 0 3 c1 c10 \n",
      "\tclose  1 4 c1 c8 \n",
      "\tcreate 0 4 c2 c9 \n",
      "\tcreate 2 4 c2 c7 \n",
      "\tclose  2 4 c2 c7 \n",
      "\tcreate 3 4 c3 c6 \n",
      "\tclose  3 4 c3 c6 \n",
      "\tcreate 3 5 c4 c7 \n",
      "\tcreate 3 3 c4 c7 \n",
      "\tcreate 4 4 c4 c5 \n",
      "\tclose  4 4 c4 c5 \n",
      "\tcreate 5 4 c5 c4 \n",
      "\tclose  5 4 c5 c4 \n",
      "\tcreate 5 5 c6 c5 \n",
      "\tcreate 6 4 c6 c3 \n",
      "\tclose  6 4 c6 c3 \n",
      "\tcreate 6 5 c7 c4 \n",
      "\tclose  0 3 c1 c10 \n",
      "\tcreate 0 2 c2 c11 \n",
      "\tclose  3 3 c4 c7 \n",
      "\tclose  0 4 c2 c9 \n",
      "\tcreate 0 5 c3 c10 \n",
      "\tclose  3 5 c4 c7 \n",
      "\tcreate 3 6 c5 c8 \n",
      "\tclose  5 5 c6 c5 \n",
      "\tclose  6 5 c7 c4 \n",
      "\tcreate 7 5 c8 c3 \n",
      "\tclose  7 5 c8 c3 \n",
      "\tcreate 7 6 c9 c4 \n",
      "\tcreate 8 5 c9 c2 \n",
      "\tclose  8 5 c9 c2 \n",
      "\tcreate 8 6 c10 c3 \n",
      "\tcreate 9 5 c10 c1 \n",
      "\tclose  9 5 c10 c1 \n",
      "\tcreate 9 6 c11 c2 \n",
      "\tcreate 9 4 c11 c0 \n",
      "\tclose  9 4 c11 c0\n",
      "Plan:            \n",
      "\tplan 1 3 \n",
      "\tplan 1 4 \n",
      "\tplan 2 4 \n",
      "\tplan 3 4 \n",
      "\tplan 4 4 \n",
      "\tplan 5 4 \n",
      "\tplan 6 4 \n",
      "\tplan 6 5 \n",
      "\tplan 7 5 \n",
      "\tplan 8 5 \n",
      "\tplan 9 5 \n",
      "\tplan 9 4\n"
     ]
    }
   ],
   "source": [
    "tok_trace = next(iter(tok_dataset.train_it(tok_dataset.train_ids)))[0]\n",
    "\n",
    "prompt_str = \" \".join(tok_trace.prompt).replace(\"start\", \"\\n\\tstart\").replace(\"goal\", \"\\n\\tgoal \").replace(\"wall\", \"\\n\\twall \")\n",
    "execution_trace_str = \" \".join(tok_trace.reasoning).replace(\"create\", \"\\n\\tcreate\").replace(\"close\", \"\\n\\tclose \")\n",
    "plan_str = \" \".join(tok_trace.plan).replace(\"plan\", \"\\n\\tplan\")\n",
    "print(f\"Prompt:          {prompt_str}\")\n",
    "print(f\"Execution trace: {execution_trace_str}\")\n",
    "print(f\"Plan:            {plan_str}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}