ReST-MCTS*

ReST-MCTS requires different environments depending on whether you are running Mistral/Llama models or SciGLM models due to transformers dependency variations. Use Miniconda to manage these environments.

• For Mistral or Llama models: Use Python 3.12.
• For SciGLM models: Use Python 3.11.

Install the required packages using the provided requirements files.

# For running Mistral (or Llama)
pip install -r requirements_mistral.txt

# For running SciGLM
pip install -r requirements_sciglm.txt

Policy data is used for training and comparing policy models across different iterations and self-training methods.

Data Characteristics:

• CoT (Chain of Thought) and MCTS datasets include only positive samples.
• DPO (Direct Preference Optimization) datasets include both positive and negative samples.

Datasets are available on Hugging Face for the following backbones:

• Llama3-8b-Instruct
• Mistral: MetaMATH-7b
• SciGLM-6B

Each backbone has data for the 1st and 2nd iterations, categorized by self-training method: ReST-EM (CoT), Self-Rewarding (DPO), and ReST-MCTS.

Process Reward Model (PRM) data, containing both positive and negative samples, is available for training the 1st reward model (specifically for Llama3-8b-Instruct).

Dataset Link: ReST-MCTS-Llama3-8b-Instruct-PRM-1st

To run MCTS* search, you must implement a policy model and a process reward model (value model). You can configure these by editing models/model.py with the following keys:

• INFERENCE_MODEL_DIR: Local path to the policy model (e.g., Llama3-8B-Instruct, Mistral-7B: MetaMATH, or SciGLM-6B).
• VALUE_BASE_MODEL_DIR: Local path to the value model backbone.
  • Use Mistral-7B if the policy is Llama3-8B-Instruct or Mistral-7B: MetaMATH.
  • Use ChatGLM3-6B if the policy is SciGLM-6B.
• VALUE_MODEL_STATE_DICT: The state dict for the value model.
• LOCAL_INFERENCE_IDX: The index corresponding to your implemented policy model.
• LOCAL_VALUE_IDX: The index corresponding to your implemented value model.

Supported implementations in models/model.py include llama, glm, and mistral for policies, and glm and mistral for value models.

You can use the MCTS_Task interface in MCTS/task.py to run a search for a specific question. The lang parameter specifies the language (e.g., 'en').

from MCTS.task import *
question = "Calculate the sum of the first 10 prime numbers."
task = MCTS_Task(question, 'llama', 'local', lang='en')
output = task.run()
print(output['solution'])

Before running MCTS* search for evaluation or generation, ensure your target question dataset is a JSON file where each item follows this format:

• content (Required): The question text.
• answer (Optional): The ground truth answer, used for evaluation.

{
  "content": "Calculate the sum of the first 10 prime numbers.",
  "answer": "129"
}

To evaluate MCTS* performance on benchmarks, use the evaluate.py script with --mode "mcts".

Key arguments:

• --task_name: The name of the benchmark task.
• --file: The specific file or subset to evaluate.
• --propose_method: The method used for proposing solutions (e.g., gpt).
• --value_method: The method used for value estimation (e.g., local).
• --mode: Set to mcts for MCTS* evaluation.
• --evaluate: The benchmark name.
• --iteration_limit: Maximum iterations.
• --use_reflection: Reflection strategy (e.g., simple).
• --branch: Branching factor for search.

python evaluate.py \
  --task_name "scibench" \
  --file "thermo" \
  --propose_method "gpt" \
  --value_method "local" \
  --mode "mcts" \
  --evaluate "scibench" \
  --iteration_limit 50 \
  --use_reflection "simple" \
  --branch 3

The self-rewarding baseline is implemented in the following script:

./self_train/self_train_dpo.py