notebooks/parse_stacks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "specialized-reception",
   "metadata": {},
   "source": [
    "# Generating the NaturalProofs Stacks domain\n",
    "\n",
    "This notebook is used to create NaturalProofs's ProoStacksfWiki domain (`naturalproofs_stacks.json`).\n",
    "\n",
    "First, pull the [Stacks github](https://github.com/stacks/stacks-project) (we used commit 4df67b8)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "divided-trinity",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import json\n",
    "import re\n",
    "from collections import defaultdict\n",
    "from tqdm import tqdm\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "dense-richards",
   "metadata": {},
   "outputs": [],
   "source": [
    "files_ = glob.glob('./stacks-project/*.tex')\n",
    "files = []\n",
    "for f in files_:\n",
    "    if 'coding.tex' in f:\n",
    "        continue\n",
    "    files.append(f)\n",
    "\n",
    "stems = [os.path.basename(f).split('.tex')[0] for f in files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "prospective-cameroon",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'definition',\n",
       " 'equation',\n",
       " 'example',\n",
       " 'exercise',\n",
       " 'item',\n",
       " 'lemma',\n",
       " 'proposition',\n",
       " 'remark',\n",
       " 'remarks',\n",
       " 'section',\n",
       " 'situation',\n",
       " 'subsection',\n",
       " 'theorem'}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_types = set()\n",
    "for f in files:\n",
    "    tex = open(f).read()\n",
    "    stem = os.path.basename(f).split('.tex')[0]\n",
    "\n",
    "    labels_ = re.findall(r'\\\\label{([a-z|A-Z|0-9|\\-]+)}', tex)\n",
    "    for l in labels_:\n",
    "        all_types.add(l.split('-')[0])\n",
    "\n",
    "all_types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceramic-level",
   "metadata": {},
   "source": [
    "### Parse dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "quality-aging",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_refs(s):\n",
    "    refs = re.findall(r'\\\\ref{([^}]*)}', s)\n",
    "    refs = [ref for ref in refs if all([t not in exclude_kinds for t in ref.split('-')])]\n",
    "    for i in range(len(refs)):\n",
    "        add_stem = True\n",
    "        for stem_ in stems:\n",
    "            if refs[i].startswith(stem_):\n",
    "                add_stem = False\n",
    "        if add_stem:\n",
    "            refs[i] = '%s-%s' % (stem, refs[i])\n",
    "    return refs\n",
    "\n",
    "def parse_proof(statement):\n",
    "    contents = statement.strip().split('\\n')\n",
    "    contents = list(filter(lambda s: s != '', contents))\n",
    "    refs = extract_refs(proof)\n",
    "    \n",
    "    return {\n",
    "        'contents': contents,\n",
    "        'refs': refs,\n",
    "    }\n",
    "\n",
    "def parse_item(statement):\n",
    "    lines = statement.strip().split('\\n')\n",
    "    start = 0\n",
    "    label = None\n",
    "    for i, line in enumerate(lines):\n",
    "        if '\\\\label' in line:\n",
    "            label = re.findall(r'\\\\label{([^}]*)}', line)[0]\n",
    "            start = i+1\n",
    "            break\n",
    "    if label is None:\n",
    "        raise ValueError('no label')\n",
    "    label = '%s-%s' % (stem, label)\n",
    "    contents = lines[start:]\n",
    "    contents = list(filter(lambda s: s != '', contents))\n",
    "    refs = extract_refs(statement)\n",
    "\n",
    "    return {\n",
    "        'label': label,\n",
    "        'categories': [stem],\n",
    "        'title': label,\n",
    "        'contents': contents,\n",
    "        'refs': refs,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "owned-queens",
   "metadata": {},
   "outputs": [],
   "source": [
    "theorem_kinds = ['theorem', 'lemma', 'proposition']\n",
    "definition_kinds = ['definition']\n",
    "other_kinds = ['remark', 'remarks']\n",
    "all_ref_kinds = theorem_kinds + definition_kinds + other_kinds\n",
    "exclude_kinds = [t for t in all_types if t not in all_ref_kinds]\n",
    "\n",
    "kind2type = {}\n",
    "for kind in theorem_kinds:\n",
    "    kind2type[kind] = 'theorem'\n",
    "for kind in definition_kinds:\n",
    "    kind2type[kind] = 'definition'\n",
    "for kind in other_kinds:\n",
    "    kind2type[kind] = 'other'\n",
    "\n",
    "theorems = []\n",
    "definitions = []\n",
    "others = []\n",
    "label2id = {}\n",
    "cnt = 0\n",
    "\n",
    "for f in files:\n",
    "    tex = open(f).read()\n",
    "    stem = os.path.basename(f).split('.tex')[0]\n",
    "\n",
    "    for kind in all_ref_kinds:\n",
    "        splits = tex.split('\\\\begin{%s}' % kind)[1:]\n",
    "        for split in splits:\n",
    "            item = {\n",
    "                'id': cnt,\n",
    "                'type': kind2type[kind],\n",
    "            }\n",
    "            cnt += 1\n",
    "            \n",
    "            statement, other = split.split('\\\\end{%s}' % kind)\n",
    "            item.update(parse_item(statement))\n",
    "\n",
    "            if kind in theorem_kinds:\n",
    "                proof = other.split('\\\\end{proof}')[0]\n",
    "                proof = re.findall(r'\\\\begin{proof}(.*)', proof, re.DOTALL)\n",
    "                assert len(proof) == 1\n",
    "                proof = proof[0]\n",
    "                proof = parse_proof(proof)\n",
    "                item['proofs'] = [proof]\n",
    "\n",
    "                theorems.append(item)\n",
    "\n",
    "            elif kind in definition_kinds:\n",
    "                definitions.append(item)\n",
    "\n",
    "            elif kind in other_kinds:\n",
    "                others.append(item)\n",
    "\n",
    "            label2id[item['label']] = item['id']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mighty-theme",
   "metadata": {},
   "source": [
    "#### Add `ref_ids`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "interested-newport",
   "metadata": {},
   "outputs": [],
   "source": [
    "for item in theorems:\n",
    "    item['ref_ids'] = [label2id[label] for label in item['refs']]\n",
    "    for proof in item['proofs']:\n",
    "        proof['ref_ids'] = [label2id[label] for label in proof['refs']]\n",
    "for item in definitions:\n",
    "    item['ref_ids'] = [label2id[label] for label in item['refs']]\n",
    "for item in others:\n",
    "    item['ref_ids'] = [label2id[label] for label in item['refs']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "challenging-mistake",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieval_examples = [thm['id'] for thm in theorems if len(thm['proofs']) > 0 and len(thm['proofs'][0]['refs']) > 0]\n",
    "\n",
    "dataset = {\n",
    "    'theorems': theorems,\n",
    "    'definitions': definitions,\n",
    "    'others': others,\n",
    "    'retrieval_examples': retrieval_examples,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "natural-valentine",
   "metadata": {},
   "source": [
    "#### Split dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "auburn-intake",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16 1-cycles\n"
     ]
    }
   ],
   "source": [
    "refs = theorems + definitions + others\n",
    "\n",
    "id2ref = {}\n",
    "label2id = {}\n",
    "for r in refs:\n",
    "    id2ref[r['id']] = r\n",
    "    label2id[r['label']] = r['id']\n",
    "    \n",
    "graph = defaultdict(list)\n",
    "pairs = []\n",
    "cycles = []\n",
    "\n",
    "for r1 in refs:\n",
    "    \n",
    "    # Make an edge for each reference in the _statement_\n",
    "    for r2 in r1['refs']:\n",
    "        \n",
    "        r1id = r1['id']\n",
    "        r2id = label2id[r2]\n",
    "        \n",
    "        if r1id != r2id:\n",
    "            graph[r2id].append(r1id)\n",
    "            \n",
    "            pairs.append((r2id, r1id))\n",
    "            \n",
    "            if r2id in graph[r1id]:\n",
    "                cycles.append(tuple(sorted((r2id, r1id))))\n",
    "                \n",
    "    # Make an edge for each reference in the _proof_ (when available)\n",
    "    if r1['type'] == 'theorem':\n",
    "        for proof in r1['proofs']:\n",
    "            \n",
    "            for r2 in proof['refs']:                \n",
    "                r1id = r1['id']\n",
    "                r2id = label2id[r2]\n",
    "                if r1id != r2id:\n",
    "                    graph[r2id].append(r1id)\n",
    "                    \n",
    "                    pairs.append((r2id, r1id))\n",
    "\n",
    "                    if r2id in graph[r1id]:\n",
    "                        cycles.append(tuple(sorted((r2id, r1id))))\n",
    "\n",
    "cycles = set(cycles)\n",
    "print(\"%d 1-cycles\" % (len(cycles)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "complete-density",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13925 nodes\n",
      "2107 leaf\n",
      "11818 non-leaf\n",
      "\n",
      "1958 heads\n"
     ]
    }
   ],
   "source": [
    "import networkx\n",
    "\n",
    "G = networkx.DiGraph(graph)\n",
    "leafs = [node for node in G.nodes() if G.in_degree(node) != 0 and G.out_degree(node)==0]\n",
    "nonleafs = [node for node in G.nodes() if G.in_degree(node) == 0 or G.out_degree(node) != 0]\n",
    "heads = [node for node in G.nodes() if G.in_degree(node) == 0 and G.out_degree(node) > 0]\n",
    "\n",
    "print(\"%d nodes\\n%d leaf\\n%d non-leaf\\n\\n%d heads\" % (\n",
    "    len(G.nodes()),\n",
    "    len(leafs),\n",
    "    len(nonleafs),\n",
    "    len(heads)\n",
    "))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cathedral-kentucky",
   "metadata": {},
   "source": [
    "#### Define the train, valid, test splits\n",
    "\n",
    "We define valid $\\cup$ test as leaves, selected at each layer proportional to the number of leaves at the layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "altered-death",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "eval_thms 1551\n",
      "train_thms 9022\n",
      "eval_refs 1551\n",
      "train_refs 13583\n"
     ]
    }
   ],
   "source": [
    "rand = np.random.RandomState(42)\n",
    "\n",
    "splits = defaultdict(set)\n",
    "\n",
    "splits['eval_thms'] = [rid for rid in retrieval_examples if rid in leafs]\n",
    "splits['train_thms'] = [rid for rid in retrieval_examples if rid not in leafs]\n",
    "splits['eval_refs'] = splits['eval_thms']\n",
    "splits['train_refs'] = [x for x in id2ref if x not in splits['eval_refs']]\n",
    "\n",
    "for k in splits:\n",
    "    splits[k] = list(splits[k])\n",
    "    print(k, len(splits[k]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hindu-aberdeen",
   "metadata": {},
   "source": [
    "#### Verify that evaluation theorems are not referred in training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "hawaiian-vision",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1551/1551 [00:04<00:00, 346.35it/s]\n"
     ]
    }
   ],
   "source": [
    "for x in tqdm(splits['eval_thms'], total=len(splits['eval_thms'])):\n",
    "    for y in splits['train_refs']:\n",
    "        if G.has_predecessor(y, x):\n",
    "            print(id2ref[x]['title'], id2ref[y]['title'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "imported-vehicle",
   "metadata": {},
   "source": [
    "#### Randomly split evaluation into validation and test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "charged-shape",
   "metadata": {},
   "outputs": [],
   "source": [
    "rand = np.random.RandomState(42)\n",
    "perm = rand.permutation(len(splits['eval_thms']))\n",
    "\n",
    "idx = len(splits['eval_thms'])//2\n",
    "val_idxs = perm[:idx]\n",
    "\n",
    "val_thms = [splits['eval_thms'][i] for i in perm[:idx]]\n",
    "tst_thms = [splits['eval_thms'][i] for i in perm[idx:]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "looking-objective",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_splits = {\n",
    "    'train': {},\n",
    "    'valid': {},\n",
    "    'test': {}\n",
    "}\n",
    "\n",
    "final_splits['train']['ref_ids'] = splits['train_refs']\n",
    "final_splits['train']['examples'] = [(tid, 0) for tid in splits['train_thms']]\n",
    "\n",
    "final_splits['valid']['ref_ids'] = splits['train_refs'] + splits['eval_refs']\n",
    "final_splits['valid']['examples'] = [(tid, 0) for tid in val_thms]\n",
    "\n",
    "final_splits['test']['ref_ids'] = splits['train_refs'] + splits['eval_refs']\n",
    "final_splits['test']['examples'] = [(tid, 0) for tid in tst_thms]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "dying-lingerie",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train\n",
      "ref_ids 13583\n",
      "examples 9022\n",
      "\n",
      "valid\n",
      "ref_ids 15134\n",
      "examples 775\n",
      "\n",
      "test\n",
      "ref_ids 15134\n",
      "examples 776\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for k in final_splits:\n",
    "    print(k)\n",
    "    for k2 in final_splits[k]:\n",
    "        print(k2, len(final_splits[k][k2]))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "going-highland",
   "metadata": {},
   "outputs": [],
   "source": [
    "js = {\n",
    "    'dataset': dataset,\n",
    "    'splits': final_splits,\n",
    "}\n",
    "\n",
    "import json\n",
    "output_json = './naturalproofs_stacks.json'\n",
    "with open(output_json, 'w') as f:\n",
    "    json.dump(js, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "basic-society",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}