ag4masses-public.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 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"
   ]
  },
  {
   "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
}