Upload 5 files

ag4masses-public.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {
     "executionInfo": {
      "elapsed": 611,
@@ -14,17 +14,16 @@
      },
      "user_tz": 300
     },
-    "id": "-IHoHd-t5sLP",
-    "trusted": true
+    "id": "-IHoHd-t5sLP"
    },
    "outputs": [],
    "source": [
     "import sys, os\n",
     "\n",
-    "AG4MDIR='/home/user/ag4masses'\n",
-    "AGLIB=f'{AG4MDIR}/aglib'\n",
-    "AGDIR=f\"{AGLIB}/alphageometry\"\n",
-    "MELIAD_PATH=f\"{AGDIR}/meliad\"\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\""
    ]
@@ -41,9 +40,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Run this cell to refresh code and get the latest versions\n",
@@ -65,8 +62,7 @@
      },
      "user_tz": 300
     },
-    "id": "GgR_vO8XX9Vr",
-    "trusted": true
+    "id": "GgR_vO8XX9Vr"
    },
    "outputs": [],
    "source": [
@@ -99,8 +95,7 @@
      "user_tz": 300
     },
     "id": "gP4zAZh2MHcv",
-    "outputId": "4796397b-8952-411e-bd33-8fd813865735",
-    "trusted": true
+    "outputId": "4796397b-8952-411e-bd33-8fd813865735"
    },
    "outputs": [],
    "source": [
@@ -147,8 +142,7 @@
      "user_tz": 300
     },
     "id": "X8Aj3G0neT6K",
-    "outputId": "9538ceba-8065-44d6-a32f-35127e5f2575",
-    "trusted": true
+    "outputId": "9538ceba-8065-44d6-a32f-35127e5f2575"
    },
    "outputs": [],
    "source": [
@@ -174,8 +168,7 @@
      "user_tz": 300
     },
     "id": "u9fuBSr2qEwN",
-    "outputId": "97bbce78-8b49-4d3b-a831-d188a4a9e536",
-    "trusted": true
+    "outputId": "97bbce78-8b49-4d3b-a831-d188a4a9e536"
    },
    "outputs": [],
    "source": [
@@ -190,9 +183,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Linux packages for Nvidia gpu.\n",
@@ -206,8 +197,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "fChy49CNhf01",
-    "trusted": true
+    "id": "fChy49CNhf01"
    },
    "outputs": [],
    "source": [
@@ -216,6 +206,51 @@
     "!nvidia-smi"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "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": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "PROB='imo-2024-q4'\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -227,10 +262,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
+   "execution_count": 15,
+   "metadata": {},
    "outputs": [],
    "source": [
     "#!! cannot have ' in the script, including in comments\n",
@@ -301,6 +334,7 @@
     "\n",
     "true \"==========================================\"\n",
     "\n",
+    "cd $AG4MDIR\n",
     "python -m alphageometry \\\n",
     "--alsologtostderr \\\n",
     "--problems_file=$PROB_FILE \\\n",
@@ -310,7 +344,7 @@
     "\"${SEARCH_ARGS[@]}\" \\\n",
     "\"${LM_ARGS[@]}\" \\\n",
     "--out_file=$OUTFILE \\\n",
-    "--n_workers=$NWORKERS 2>&1\n",
+    "--n_workers=$NWORKERS # 2>&1\n",
     "\n",
     "'''"
    ]
@@ -318,10 +352,18 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
-   "outputs": [],
+   "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"
+     ]
+    }
+   ],
    "source": [
     "os.environ[\"TESTDIR\"]=TESTDIR\n",
     "os.environ[\"AG4MDIR\"]=AG4MDIR\n",
@@ -341,16 +383,16 @@
     "#   NWORKERS=2\n",
     "#   CUDA_VISIBLE_DEVICES=0,1\n",
     "\n",
-    "os.environ[\"BATCH_SIZE\"]=\"16\"\n",
-    "os.environ[\"BEAM_SIZE\"]=\"64\"\n",
-    "os.environ[\"DEPTH\"]=\"8\"\n",
+    "os.environ[\"BATCH_SIZE\"]=\"2\"\n",
+    "os.environ[\"BEAM_SIZE\"]=\"2\"\n",
+    "os.environ[\"DEPTH\"]=\"2\"\n",
     "os.environ[\"NWORKERS\"]=\"2\"\n",
     "\n",
-    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0,1\"\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[\"PROB_FILE\"]=f\"{AG4MDIR}/data/ag4m_problems.txt\"\n",
-    "PROB=\"imo-2024-q4\"\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",
@@ -358,18 +400,15 @@
     "# 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",
-    "else:\n",
-    "    os.system(f\"echo '{jobScript}'|bash &\")\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": {
-    "trusted": true
-   },
+   "execution_count": 6,
+   "metadata": {},
    "outputs": [],
    "source": [
     "#!cat /kaggle/input/tmpfiles/test-problems.txt"
@@ -378,9 +417,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# In an interactive Kaggle session, run this to see the log file. We can cancel this cell's execution\n",
@@ -392,9 +429,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Command to kill the background job in an interactive session\n",
@@ -409,9 +444,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Command to check progress of a running job in an interactive session\n",
@@ -421,9 +454,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "trusted": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# In Batch run, after the job completes, list output files\n",
@@ -451,7 +482,7 @@
    "sourceType": "notebook"
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -465,7 +496,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

download.sh

@@ -7,7 +7,9 @@ git clone https://github.com/tpgh24/ag4masses.git
 pip cache purge
 pip install --upgrade pip
 pip install --upgrade packaging setuptools setuptools_scm wheel 
+# pip install typing_extensions==4.6.0
 pip install --require-hashes --no-deps -r /home/user/app/aglib/ag4masses/alphageometry/requirements.txt
+pip install typing_extensions==4.6.0
 
 # cd alphageometry
 git clone https://github.com/google-research/meliad.git
@@ -21,5 +23,5 @@ export DATA=ag_ckpt_vocab
 cd /home/user/app
 # some patch for cpu
 cp models.py /home/user/app/aglib/ag4masses/alphageometry/models.py
-cp alphageometry.py /home/user/app/aglib/ag4masses/alphageometry/alphageometry.py
+cp lm_inference.py /home/user/app/aglib/ag4masses/alphageometry/lm_inference.py
 cp ag4masses-public.ipynb /data/ag4masses-public.ipynb

lm_inference.py

@@ -0,0 +1,189 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Wrapper for language modeling inference implemented in Meliad."""
+from typing import Any, Dict
+
+import jax
+import models  # pylint: disable=unused-import
+import t5.data
+from transformer import inference_utils
+
+
+np = jax.numpy
+
+
+Trainer = inference_utils.Trainer
+
+MetricsOutput = Dict[str, Any]  # Metrics output by model.
+
+
+parse_gin_configuration = inference_utils.parse_gin_configuration
+
+
+class LanguageModelInference:
+  """Meliad wrapper for LM inference."""
+
+  def __init__(self, vocab_path: str, load_dir: str, mode='beam_search'):
+    self.vocab = t5.data.SentencePieceVocabulary(vocab_path)
+
+    # This task won't be pulling from a dataset.
+    def null_iter_fn() -> None:
+      return None
+
+    process_summaries_f = inference_utils.models.process_summaries_function(
+        self.vocab
+    )
+
+    trainer = inference_utils.training_loop.Trainer(
+        get_training_dataset_iterator=null_iter_fn,
+        get_test_dataset_iterator=None,
+        pretty_print_input_function=None,
+        process_summaries_function=process_summaries_f,
+        load_dir=load_dir,
+        workdir='',  # Don't log or save checkpoints.
+        replicate_mode=False,
+    )  # Run on a single device at batch size 1.
+    self.trainer = trainer
+
+    # Create and initialize the model.
+    (tstate, _, imodel, prngs) = trainer.initialize_model()
+    self.imodel = imodel
+    self.batch_size = imodel.task_config.batch_size
+
+    self.n = imodel.num_heads
+    self.h = imodel.head_size
+
+    # Create an inference task.
+    writers = {}
+    self.task = trainer.create_training_task(mode, imodel, prngs, writers)  # pylint: disable=too-many-function-args
+
+    # Register any additional actions.
+    # Actions are cleared first for use with colab.
+    inference_utils.training_loop.clear_interstep_callbacks()
+    inference_utils.training_loop.register_interstep_callbacks()
+    self.tstate = tstate
+
+    # some default parameters.
+    eos = [0] * 1024
+    for idx in self.encode_list(['.', ';']):
+      eos[idx] = 1
+
+    self.eos = np.array(eos, dtype=np.float32)
+    self.mask = jax.numpy.ones([1024], dtype=np.float32)
+
+  def decode(self, ids: list[int]) -> str:
+    return self.vocab.decode(ids)
+
+  def decode_list(self, tokens: list[int]) -> list[str]:
+    return [self.decode([tok]) for tok in tokens]
+
+  def encode(self, inputs_str: str) -> list[int]:
+    return self.vocab.encode(inputs_str)
+
+  def encode_list(self, inputs_strs: list[str]) -> list[int]:
+    result = [self.vocab.encode(x) for x in inputs_strs]
+    assert all([len(x) == 1 for x in result]), [
+        self.decode(x) for x in result if len(x) != 1
+    ]
+    return [x[0] for x in result]
+
+  def call(
+      self,
+      inputs: np.ndarray,
+      dstate: tuple[dict[str, np.ndarray], ...] = None,
+      eos: np.ndarray = None,
+      mask: np.ndarray = None,
+  ) -> MetricsOutput:
+    """Call the meliad model."""
+    batch_size, length = inputs.shape
+    inputs = jax.numpy.pad(inputs, [(0, 0), (0, 1024 - length)])
+
+    if eos is None:
+      eos = self.eos
+    if mask is None:
+      mask = self.mask
+
+    x = {'targets': inputs, 'length': length, 'eos': eos, 'mask': mask}
+
+    if dstate is not None:
+      x['start_of_sequence'] = jax.numpy.array([False] * batch_size)
+    else:
+      dstate = tuple(
+          [{  # this dummy value will never be used.
+              'current_index': np.array([0] * batch_size, dtype=np.int32),
+              'keys': np.zeros(
+                  (batch_size, 2048, self.n, self.h), dtype=np.float32
+              ),
+              'values': np.zeros(
+                  (batch_size, 2048, self.n, self.h), dtype=np.float32
+              ),
+              'recurrent_kvq': None,
+              'relative_position_bias': np.zeros(
+                  (batch_size, self.n, 1, 1024), dtype=np.float32
+              ),
+          }]
+          * 12
+      )
+      x['start_of_sequence'] = jax.numpy.array([True] * batch_size)
+
+    x['dstate'] = dstate
+    _, metrics_np = self.task.run_step(self.tstate, x, 0)
+    return metrics_np
+
+  def beam_decode(
+      self,
+      inputs: str,
+      eos_tokens: np.ndarray = None,
+      mask_tokens: np.ndarray = None,
+      dstate: dict[str, np.ndarray] = None,
+  ) -> MetricsOutput:
+    """Beam search."""
+    inputs = jax.numpy.array([self.vocab.encode(inputs)] * self.batch_size)
+
+    eos = self.eos
+    if eos_tokens is not None:
+      eos_ids = self.encode_list(eos_tokens)
+      eos = np.array(
+          [1 if idx in eos_ids else 0 for idx in range(1024)], dtype=np.float32
+      ).reshape((1, 1, 1024))
+
+    mask = self.mask
+    if mask_tokens is not None:
+      mask_ids = self.encode_list(mask_tokens)
+      mask = np.array(
+          [0 if idx in mask_ids else 1 for idx in range(1024)],
+          dtype=np.float32,
+      ).reshape((1, 1, 1024))
+
+    metrics_np = self.call(inputs, dstate=dstate, eos=eos, mask=mask)
+
+    finished_seqs = metrics_np['finished_seqs']
+    finished_scores = metrics_np['finished_scores']
+
+    seqs = []
+    scores = []
+    for seq, score in zip(finished_seqs, finished_scores):
+      seq = self.decode(seq[1:])
+      seqs.append(seq)
+      scores.append(score)
+
+    return {
+        'finished_seqs': finished_seqs,
+        'finished_scores': finished_scores,
+        'seqs_str': seqs,
+        'scores': scores,
+        'dstate': metrics_np['dstate'],
+    }