!git clone https://github.com/huggingface/search-and-learn

%cd search-and-learn
!pip install -e '.[dev]'

from huggingface_hub import notebook_login

notebook_login()

import torch
from vllm import LLM
from sal.models.reward_models import RLHFFlow

model_path = "meta-llama/Llama-3.2-1B-Instruct"
prm_path = "RLHFlow/Llama3.1-8B-PRM-Deepseek-Data"

llm = LLM(
    model=model_path,
    gpu_memory_utilization=0.5,  # Utilize 50% of GPU memory
    enable_prefix_caching=True,  # Optimize repeated prefix computations
    seed=42,  # Set seed for reproducibility
)

prm = RLHFFlow(prm_path)

question_text = "Convert the point $(0,3)$ in rectangular coordinates to polar coordinates.  Enter your answer in the form $(r,\theta),$ where $r > 0$ and $0 \le \theta < 2 \pi.$"
input_batch = {"problem": [question_text]}

from sal.config import Config
from sal.search import beam_search, best_of_n, dvts

config = Config()
config.n = 32  # Number of answers to generate during the search

search_result = best_of_n(x=input_batch, config=config, llm=llm, prm=prm)

search_result["pred"][0]

formatted_output = search_result["pred"][0].replace("<|start_header_id|>assistant<|end_header_id|>\n\n", "").strip()
formatted_output

from IPython.display import display, Markdown

display(Markdown(formatted_output))

import time


def generate_with_search_and_learn(question, config, llm, prm, method="best_of_n"):
    """
    Generate an answer for a given question using the search-and-learn pipeline.

    Args:
    - question (str): The input question to generate an answer for.
    - config (Config): Configuration object containing parameters for search strategy.
    - llm (LLM): Pretrained large language model used for generating answers.
    - prm (RLHFFlow): Process reward model used for evaluating answers.
    - method (str): Search strategy to use. Options are 'best_of_n', 'beam_search', 'dvts'. Default is 'best_of_n'.

    Returns:
    - str: The formatted output after processing the question.
    """
    batch = {"problem": [question]}

    start_time = time.time()
    if method == "best_of_n":
        result = best_of_n(x=batch, config=config, llm=llm, prm=prm)
    elif method == "beam_search":
        result = beam_search(examples=batch, config=config, llm=llm, prm=prm)
    elif method == "dvts":
        result = dvts(examples=batch, config=config, llm=llm, prm=prm)

    elapsed_time = time.time() - start_time
    print(f"\nFinished in {elapsed_time:.2f} seconds\n")

    tokenizer = llm.get_tokenizer()
    total_tokens = 0
    for completion in result["completions"]:
        for comp in completion:
            output_tokens = tokenizer.encode(comp)
            total_tokens += len(output_tokens)

    print(f"Total tokens in all completions: {total_tokens}")

    formatted_output = result["pred"][0].replace("<|start_header_id|>assistant<|end_header_id|>\n\n", "").strip()
    return formatted_output

>>> question = "Convert the point $(0,3)$ in rectangular coordinates to polar coordinates.  Enter your answer in the form $(r,\theta),$ where $r > 0$ and $0 \le \theta < 2 \pi.$"

>>> config.n = 8

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="best_of_n"
... )

display(Markdown(formatted_output))

>>> config.n = 8
>>> # beam search specific
>>> config.sort_completed = True
>>> config.filter_duplicates = True

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="beam_search"
... )

display(Markdown(formatted_output))

>>> config.n = 8
>>> # dvts specific
>>> config.n_beams = config.n // config.beam_width

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="dvts"
... )

display(Markdown(formatted_output))

>>> question = "What's the capital of Spain?"

>>> config.n = 32

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="best_of_n"
... )

display(Markdown(formatted_output))