Using kluster.ai for batch inference
Last updated
export KLUSTERAI_API_KEY=<your_api_key>"""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"],
}
]
reasoner = Reasoner(model_name="deepseek-ai/DeepSeek-R1",
backend="klusterai",
batch=True,
backend_params={"max_retries": 1, "completion_window": "1h"})from datasets import load_dataset
dataset = load_dataset("openai/gsm8k", name="main")
dataset_to_use = dataset["train"].take(3)
output = reasoner(dataset).datasetfrom 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>"
))