{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 不要保存此文件！！！ \n",
    "# AlphaGeometry\n",
    "由DeepMind开源的AlphaGeometry用于几何解题工具。\n",
    "\n",
    "## 一.使用方法\n",
    "\n",
    "### 1. 上传题目\n",
    "\n",
    "双击左侧problems.txt，在末尾换行后添加新的题目，格式见第二部分。该文件已经有部分例子\n",
    "\n",
    "### 2. 修改配置\n",
    "\n",
    "在下方代码块中直接修改PROB的值，修改为题目名称。\n",
    "\n",
    "### 3. 运行\n",
    "\n",
    "从上之下依次点击代码块左侧的运行按钮即可，或者点击上方的双箭头按钮运行全部代码块。\n",
    "\n",
    "### 4. 查看结果\n",
    "\n",
    "运行结束后，双击打开左侧的ag4mtest文件夹，双击打开`题目名.out`文件。\n",
    "\n",
    "## 二.题目格式\n",
    "\n",
    "在problems.txt中新建一行，写上题目名字，如`test`。\n",
    "\n",
    "再新建一行，开始写题目具体内容。\n",
    "\n",
    "文件中已有部分示例可以参考。\n",
    "\n",
    "下面是英文文档，懒得翻译了（如果看不懂可以让我帮忙写或szmh私信轰炸催我翻译）\n",
    "\n",
    "### The Problem Definition Language\n",
    "\n",
    "Below is a problem:\n",
    "\n",
    "```\n",
    "orthocenter\n",
    "a b c = triangle; h = on_tline b a c, on_tline c a b ? perp a h b c\n",
    "```\n",
    "\n",
    "* A problem consists of 2 lines, the first line is the name of the problem, the second line is the definition\n",
    "* The problem definition is **sensitive to white spaces, including trailing ones**\n",
    "* The problem definition consists of premises and a conclusion, separated by `' ? '`\n",
    "* The premises consist of multiple clauses for constructing points, the best way to understand them is to think of the process of drawing the points one by one\n",
    "* Multiple point-construction clauses are separated by `' ; '`. Note that the last one should **not** end with `' ; '`, before the `' ? '` separating the premises and the conclusion\n",
    "* Some point-construction clauses can construct multiple points, such as `'a b c = triangle'`\n",
    "* A point-construction clause consists of point names (separated by a single space), followed by `' = '`, and 1 or 2 \"actions\" (the term used in the Google paper), separated by `' , '`. See in the above example: `h = on_tline b a c, on_tline c a b`\n",
    "* Actions are defined in the `alphageometry/defs.txt` file. They are also listed in the Google paper in *\"Extended Data Table 1 | List of actions to construct the random premises\"* (reproduced [here](https://github.com/tpgh24/ag4masses/blob/main/data/ag_defs.jpg)). Each action is a constraint on the position of the point. Constructing a point using actions is similar to constructing it using straight edge and compass, *e.g.* find the point through intersection of 2 lines\n",
    "* An action is similar to a function call, with other points being inputs and the point to be constructed being output\n",
    "* Output point names can be optionally repeated in the beginning of the inputs (arguments) of the actions. For example, `h = on_tline b a c, on_tline c a b` can also be `h = on_tline h b a c, on_tline h c a b`. In `alphageometry/defs.txt` the output point names are repeated in front of the input point names. This sometimes makes the action clearer to read\n",
    "* It's possible to add actions but it's not enough to just add into the `defs.txt` file. In `defs.txt`, each action is defined by 5 lines. The last line invoves functions needed for numerical checking that need to be implemented in Python\n",
    "* The conclusion (goal) part of the problem can have one of the following statements:\n",
    "  * `coll a b c` : points `a b c` are collinear\n",
    "  * `cong a b c e` : segments `ab` and `cd` are congruent (length equal)\n",
    "  * `contri a b c p q r` : triangles `abc` and `pqr` are congruent\n",
    "  * `cyclic a b c d` : 4 points `a b c d` are cocyclic\n",
    "  * `eqangle a b c d p q r s` : the angles between lines `ab-cd` and `pq-rs` are equal. **Note that angles have directions (signs)** so the order between `a b` and `c d` matters. `eqangle a b c d c d a b` is false. The way to think about it is, angle `ab-cd` is the angle to turn line `ab` **clockwise** so it is parallel with the line `cd`. You can use counter-clockwise as the convention too, as long as for all angles the same convention is used\n",
    "  * `eqratio a b c d p q r s` : segment length `ab/cd = pq/rs`\n",
    "  * `midp m a b` : point `m` is the midpoint of `a` and `b`\n",
    "  * `para a b c d` : segments `ab` and `cd` are parallel\n",
    "  * `perp a b c d` : segments `ab` and `cd` are perpendicular to each other\n",
    "  * `simtri a b c p q r` : triangles `abc` and `pqr` are similar\n",
    "\n",
    "### Some Tips\n",
    "\n",
    "* **Angles have directions (signs)**. See the note for `eqangle` above. Attention needs to be paid both in the premise (point construction) part and the conclusion part of a problem\n",
    "\n",
    "* AlphaGeometry does not do robust error checking of the problem or the proposition. If the problem has syntax errors or the proposition is false, it often freezes. To detect this, look at the log on stderr. AlphaGeometry will first try to solve the problem using DD+AR, and on stderr, you should see logs like this:\n",
    "\n",
    "```\n",
    "I0324 19:53:37.293019 123295230480384 graph.py:498] pascal\n",
    "I0324 19:53:37.293379 123295230480384 graph.py:499] a = free a; b = free b; c = on_circle c a b; d = on_circle d a b; e = on_circle e a b; f = on_circle f a b; g = on_circle g a b; h = intersection_ll h b c e f; i = intersection_ll i c d f g; j = intersection_ll j d e g b ? coll h i j\n",
    "I0324 19:53:38.638956 123295230480384 ddar.py:60] Depth 1/1000 time = 1.2907805442810059\n",
    "I0324 19:53:42.962377 123295230480384 ddar.py:60] Depth 2/1000 time = 4.3230626583099365\n",
    "I0324 19:53:47.302527 123295230480384 ddar.py:60] Depth 3/1000 time = 4.3398051261901855\n",
    "```\n",
    "\n",
    "Using the AG4Masses code, this should happen right away. Using the original AlphaGeometry code, when the model is `alphageometry`, it will take several minutes to get there because the original AlphaGeometry code loads the LM first. In any case, if you do not see this after several minutes, chances are there is an error in the syntax of the problem or the proposition is false.\n",
    "\n",
    "One trick to error-check a problem's syntax and generate the diagram for the problem is to first use a trivial conclusion such as `cong a b a b`. If the rest of the problem is correct, it will be proven right away, and you will get a diagram generated by the code. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "PROB='imo-2024-q4'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "executionInfo": {
     "elapsed": 611,
     "status": "ok",
     "timestamp": 1733595497864,
     "user": {
      "displayName": "Tong Peng",
      "userId": "14680520704856526492"
     },
     "user_tz": 300
    },
    "id": "-IHoHd-t5sLP"
   },
   "outputs": [],
   "source": [
    "import sys, os\n",
    "\n",
    "AG4MDIR='/home/user/app/aglib/ag4masses'\n",
    "AGLIB=f'/home/user/app/aglib/'\n",
    "AGDIR=f\"{AG4MDIR}/alphageometry\"\n",
    "MELIAD_PATH=f\"{AGLIB}/meliad\"\n",
    "DATA=f\"{AGLIB}/ag_ckpt_vocab\"\n",
    "TESTDIR=f\"/data/ag4mtest\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jUWvch7kYhxt"
   },
   "source": [
    "# Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!! cannot have ' in the script, including in comments\n",
    "jobScript = '''\n",
    "# !/bin/bash\n",
    "set -e\n",
    "set -x\n",
    "\n",
    "# stdout, solution is written here\n",
    "OUTFILE=$TESTDIR/${PROB}.out\n",
    "# stderr, a lot of information, error message, log etc.\n",
    "ERRFILE=$TESTDIR/${PROB}.log\n",
    "\n",
    "# stdout and stderr are written to both ERRFILF and console\n",
    "exec >$ERRFILE 2>&1\n",
    "\n",
    "echo PROB=$PROB\n",
    "echo PROB_FILE=$PROBFILE\n",
    "echo MODEL=$MODEL\n",
    "\n",
    "# Directory where output files go\n",
    "echo TESTDIR=$TESTDIR\n",
    "# Directory containing AG4Masses source files\n",
    "echo AG4MDIR=$AG4MDIR\n",
    "# Directory containing external libraries including ag_ckpt_vocab and meliad\n",
    "echo AGLIB=$AGLIB\n",
    "\n",
    "AGDIR=$AG4MDIR/alphageometry\n",
    "export PYTHONPATH=$PYTHONPATH:$AGDIR:$AGLIB\n",
    "\n",
    "echo BATCH_SIZE=$BATCH_SIZE\n",
    "echo BEAM_SIZE=$BEAM_SIZE\n",
    "echo DEPTH=$DEPTH\n",
    "echo NWORKERS=$NWORKERS\n",
    "\n",
    "echo ERRFILE=$ERRFILE\n",
    "echo OUTFILE=$OUTFILE\n",
    "\n",
    "DATA=$AGLIB/ag_ckpt_vocab\n",
    "MELIAD_PATH=$AGLIB/meliad\n",
    "export PYTHONPATH=$PYTHONPATH:$MELIAD_PATH\n",
    "\n",
    "DDAR_ARGS=( \\\n",
    "  --defs_file=$AGDIR/defs.txt \\\n",
    "  --rules_file=$AGDIR/rules.txt \\\n",
    ")\n",
    "\n",
    "SEARCH_ARGS=(\n",
    "  --beam_size=$BEAM_SIZE\n",
    "  --search_depth=$DEPTH\n",
    ")\n",
    "\n",
    "LM_ARGS=(\n",
    "  --ckpt_path=$DATA \\\n",
    "  --vocab_path=$DATA/geometry.757.model \\\n",
    "  --gin_search_paths=$MELIAD_PATH/transformer/configs,$AGDIR \\\n",
    "  --gin_file=base_htrans.gin \\\n",
    "  --gin_file=size/medium_150M.gin \\\n",
    "  --gin_file=options/positions_t5.gin \\\n",
    "  --gin_file=options/lr_cosine_decay.gin \\\n",
    "  --gin_file=options/seq_1024_nocache.gin \\\n",
    "  --gin_file=geometry_150M_generate.gin \\\n",
    "  --gin_param=DecoderOnlyLanguageModelGenerate.output_token_losses=True \\\n",
    "  --gin_param=TransformerTaskConfig.batch_size=$BATCH_SIZE \\\n",
    "  --gin_param=TransformerTaskConfig.sequence_length=128 \\\n",
    "  --gin_param=Trainer.restore_state_variables=False\n",
    ");\n",
    "\n",
    "true \"==========================================\"\n",
    "\n",
    "cd $AG4MDIR\n",
    "python -m alphageometry \\\n",
    "--alsologtostderr \\\n",
    "--problems_file=$PROBFILE \\\n",
    "--problem_name=$PROB \\\n",
    "--mode=$MODEL \\\n",
    "\"${DDAR_ARGS[@]}\" \\\n",
    "\"${SEARCH_ARGS[@]}\" \\\n",
    "\"${LM_ARGS[@]}\" \\\n",
    "--out_file=$OUTFILE \\\n",
    "--n_workers=$NWORKERS 2>&1\n",
    "\n",
    "echo =======================================\n",
    "echo Task Done.\n",
    "echo See ag4mtest/$PROB.out and ag4mtest/$PROB.log for more information\n",
    "\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "+ OUTFILE=/data/ag4mtest/imo-2024-q4.out\n",
      "+ ERRFILE=/data/ag4mtest/imo-2024-q4.log\n",
      "+ exec\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "256"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.environ[\"TESTDIR\"]=TESTDIR\n",
    "os.environ[\"AG4MDIR\"]=AG4MDIR\n",
    "os.environ[\"AGLIB\"]=AGLIB\n",
    "\n",
    "# BATCH_SIZE: number of outputs for each LM query\n",
    "# BEAM_SIZE: size of the breadth-first search queue\n",
    "# DEPTH: search depth (number of auxilary points to add)\n",
    "# NWORKERS: number of parallel run worker processes.\n",
    "# \n",
    "# Memory usage is affected by BATCH_SIZE, NWORKER and complexity of the problem.\n",
    "# Larger NWORKER and BATCH_SIZE tends to cause out of memory issue\n",
    "#\n",
    "# The results in Google paper can be obtained by setting BATCH_SIZE=32, BEAM_SIZE=512, DEPTH=16\n",
    "#\n",
    "# 1/2025: Kaggle free version provides GPU T4x2, 4 virtual CPUs, 29G RAM. Can set \n",
    "#   NWORKERS=2\n",
    "#   CUDA_VISIBLE_DEVICES=0,1\n",
    "\n",
    "os.environ[\"BATCH_SIZE\"]=\"2\"\n",
    "os.environ[\"BEAM_SIZE\"]=\"2\"\n",
    "os.environ[\"DEPTH\"]=\"2\"\n",
    "os.environ[\"NWORKERS\"]=\"2\"\n",
    "\n",
    "# o# s.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0,1\"\n",
    "\n",
    "# test problems can be uploaded into a dataset, e.g. for dataset \"tmpfiles\", \"/kaggle/input/tmpfiles/test-problems.txt\"\n",
    "os.environ[\"PROBFILE\"]=\"/data/problems.txt\"\n",
    "# PROB=\"imo-2024-q4\"\n",
    "os.environ[\"PROB\"]=PROB\n",
    "# alphageometry|ddar\n",
    "os.environ[\"MODEL\"]=\"alphageometry\"\n",
    "\n",
    "# In an interactive Kaggle session, run the job in background, so we can do other things in the Notebook.\n",
    "# For long jobs, commit the Notebook and run in Batch mode.\n",
    "# An interactive session will be terminated after about 20 minutes of idle time.\n",
    "# if os.environ[\"KAGGLE_KERNEL_RUN_TYPE\"]==\"Batch\":\n",
    "os.system(f\"echo '{jobScript}'|bash\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyOcsgkfOgCk5oTpUiS6zrgo",
   "collapsed_sections": [
    "pW2KIijZBAdh"
   ],
   "gpuType": "T4",
   "provenance": []
  },
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [],
   "dockerImageVersionId": 30823,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 4
}