llm-mcts-inference 0.1.4
An experimental project using Monte Carlo Tree Search (MCTS) to refine LLM responses for better accuracy and decision-making.
Requires Python: >=3.11

# LLM MCTS Inference

![Tests](https://img.shields.io/github/actions/workflow/status/brotSchimmelt/LLM-MCTS-Inference/test.yml?style=flat-square&label=Tests&color=green)
![Publish](https://img.shields.io/github/actions/workflow/status/brotSchimmelt/LLM-MCTS-Inference/publish.yml?style=flat-square&label=Publish&color=blue)
![PyPI](https://img.shields.io/pypi/v/llm-mcts-inference)

An experimental project using Monte Carlo Tree Search (MCTS) to refine Language Model (LLM) responses for better accuracy and decision-making.

## Overview

This project leverages MCTS to explore multiple answer candidates generated by an LLM. By iteratively generating an initial answer, evaluating it, and refining it based on targeted self-feedback, the system strives to improve response quality and decision-making. This approach leverages test-time compute to enhance the precision and robustness of model outputs.

## MCTS Inference Process

The process follows these key steps:

- **Initial Answer Generation:** Uses greedy decoding to generate an initial response.
- **Feedback Generation:** Provides constructive, concise feedback on initial answers. The feedback is generated by the model itself.
- **Iterative Refinement:** Refines responses based on the feedback through additional model queries.
- **Monte Carlo Tree Search:** Employs MCTS to explore and evaluate multiple answer paths.

## Experimental Results

The performance of this approach was evaluated on a subset of the [GSM8k](https://huggingface.co/datasets/openai/gsm8k) test split using the Llama3.2-1B-instruct model with [vLLM](https://github.com/vllm-project/vllm). A baseline run using zero-shot prompting achieved a pass@8 score of **74%** and a majority@8 score of **27%**. When applying MCTS for iterative refinement, the pass@8 score marginally increased to **75%**, while the majority@8 score improved significantly to **39%**. The evaluation was done with [llm-eval](https://github.com/brotSchimmelt/llm-eval).

These results suggest that while MCTS does not drastically improve the probability of generating at least one correct answer (pass@8), it significantly enhances response consistency (majority@8), making the model more reliable in decision-making scenarios.

### Why Llama3.2-1B-instruct?

A smaller model was selected for this experiment to better illustrate the impact of MCTS. Larger models already achieve high accuracy on GSM8k, making it difficult to demonstrate meaningful improvements. The 1B parameter model provides a more realistic proof-of-concept by:
 • Being resource-efficient, allowing for scalable experimentation.
 • Providing a challenging test case, as smaller models struggle more with GSM8k, making improvements more noticeable.
 • Ensuring the evaluation remains relevant, since GSM8k has been extensively benchmarked by larger models, leaving little room for additional gains.

## Installation

### Dependencies

- **Python:** Version 3.11 or higher

The project depends mainly on the following packages:

- `instructor` for guided generation
- `litellm` provides a unified API to interact with multiple LLM providers

### Setup Instructions

To install the package directly from PyPi run the following command: `pip install llm-mcts-inference`

To install from source, follow these steps:

1. **Clone the Repository:**

   ```bash
   git clone https://github.com/brotSchimmelt/llm-mcts-inference.git
   cd llm-mcts-inference
   ```

2. **Install the Project Dependencies**:

    If you use `uv`, run the following commands to create a virtualenv and install all requirements:

    ```bash
    uv venv --python 3.11
    uv sync
    ```

    Otherwise, install the required packages with pip:

    ```bash
    pip install -r pyproject.toml
    ```

3. **Configure Environment Variables**:
    Rename the provided example.env file to .env and update it with your API keys or other configuration details as needed.

## Usage

Use the MonteCarloLLM class to generate and improve responses via MCTS:

```python
from llm_mcts_inference.MonteCarloLLM import MonteCarloLLM

# Initialize with a specific model; defaults are defined in settings
llm = MonteCarloLLM(model_name="openai/gpt-4o-mini")

# Define your prompt
prompt = "What is the capital of France?"

# Generate a response using Monte Carlo Tree Search
result = llm.generate(prompt=prompt, iterations=5, max_children=3)

# Output the final improved answer
print("Final Answer:", result.answer)

# Optionally, display the sequence of nodes (answers) along the best path
print("Best Path:", [node.answer for node in result.valid_path])
```

## License

This project is licensed under the MIT license.