Using kluster.ai for batch inference

You can use kluster.ai for batch inference in Curator to generate synthetic data. In this example, we will generate answers for GSM8K dataset, but the approach can be adapted for any data generation task.

Prerequisites

• Python 3.10+

• Curator: Install via pip install bespokelabs-curator

• kluster.ai API key: Get your key from https://www.kluster.ai/

Steps

1. Setup environment vars

export KLUSTERAI_API_KEY=<your_api_key>

2. Create a curator.LLM subclass

Create a class that inherits from curator.LLM. Implement two key methods:

• prompt(): Generates the prompt for the LLM.

• parse(): Processes the LLM's response into your desired format.

Here’s the implementation:

"""Example of reannotating the WildChat dataset using curator."""

import logging
from bespokelabs import curator

# To see more detail about how batches are being processed
logger = logging.getLogger("bespokelabs.curator")
logger.setLevel(logging.INFO)

class Reasoner(curator.LLM):
    """Curator class for processing GSM8K dataset."""

    def prompt(self, input):
        """Create a prompt for the LLM to reason about the problem."""
        return f"Answer the following question: {input['question']}"

    def parse(self, input, response):
        """Parse the LLM response to extract reasoning and solution.

        The response format is expected to be '<think>reasoning</think>answer'
        """
        full_response = response

        # Extract reasoning and answer using regex
        import re

        reasoning_pattern = r"<think>(.*?)</think>"
        reasoning_match = re.search(reasoning_pattern, full_response, re.DOTALL)

        reasoning = reasoning_match.group(1).strip() if reasoning_match else ""
        # Answer is everything after </think>
        answer = re.sub(reasoning_pattern, "", full_response, flags=re.DOTALL).strip()

        return [
            {
                "question": input["question"],
                "reasoning": reasoning,
                "deepseek_solution": answer,
                "gold_answer": input["answer"],
            }
        ]

3. Configure Reasoner to use DeepSeek-R1 through kluster.ai

reasoner = Reasoner(model_name="deepseek-ai/DeepSeek-R1", 
                    backend="klusterai", 
                    batch=True, 
                    backend_params={"max_retries": 1, "completion_window": "1h"})

4 Generate Data

Generate the structured data and output the results as a pandas DataFrame:

from datasets import load_dataset

dataset = load_dataset("openai/gsm8k", name="main")
dataset_to_use = dataset["train"].take(3)
output = reasoner(dataset)

Example Output

Using the above example, the output might look like this:

from IPython.display import HTML, display, Markdown
which = 0
question = output[which]['question']
gold_answer = output[which]['gold_answer']
model_answer = output[which]['deepseek_solution']
thought = output[which]['reasoning']

to_display_input = question.replace("\n", "<br>")
to_display_output = model_answer.replace("\n", "<br>")

display(Markdown(
    "<h1>Question</h1>"
    f"<h3>{question}</h3>"
))
display(Markdown(
    "<h1>Model answer</h1>"
    f"<p>{model_answer}</p>"
))
display(Markdown(
    "<h1>Gold answer</h1>"
    f"<p>{gold_answer}</p>"
))
display(Markdown(
    "<h1>Model Thought</h1>"
    f"<p>{thought}</p>"
))

Batch Configuration

• Check out complete batch configuration