mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/examples/anymaths-bench/eval_default.py

@@ -0,0 +1,111 @@
+from train_anymaths import init_dataset
+
+from mantisdk.algorithm.gepa.lib.adapters.anymaths_adapter.anymaths_adapter import AnyMathsStructuredOutput
+
+if __name__ == "__main__":
+    import argparse
+    import ast
+    from pathlib import Path
+
+    import litellm
+    from tqdm import tqdm
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--anymaths_dset_name", type=str, default="openai/gsm8k")
+    parser.add_argument("--model", type=str, default="ollama/qwen3:4b", help="The model to evaluate.")
+    parser.add_argument("--use_api_url", action="store_true", help="Whether to use the API URL.")
+    parser.add_argument("--api_url", type=str, default="http://localhost:11434", help="The API URL to use.")
+    parser.add_argument("--batch_size", type=int, default=8, help="The batch size for evaluation.")
+    parser.add_argument(
+        "--max_litellm_workers", type=int, default=1, help="The maximum number of LiteLLM workers to use."
+    )
+    parser.add_argument(
+        "--which_prompt",
+        type=str,
+        default="seed",
+        choices=["seed", "optimized"],
+        help="The prompt to use for evaluation.",
+    )
+
+    args = parser.parse_args()
+
+    dataset = args.anymaths_dset_name
+
+    use_api_url = args.use_api_url
+    if not use_api_url:
+        api_url = ""
+    else:
+        api_url = args.api_url
+
+    model = args.model
+    max_litellm_workers = args.max_litellm_workers
+
+    _, _, testset = init_dataset(dataset)
+
+    if args.which_prompt == "seed":
+        INSTRUCTION_PROMPT_PATH = Path(__file__).parent / "prompt-templates/instruction_prompt.txt"
+    else:
+        INSTRUCTION_PROMPT_PATH = Path(__file__).parent / "prompt-templates/optimal_prompt.txt"
+
+    instruction = INSTRUCTION_PROMPT_PATH.read_text()
+
+    batched_testset = []
+    batch_size = args.batch_size
+
+    for i in range(0, len(testset), batch_size):
+        batched_testset.append(testset[i : i + batch_size])
+
+    total_score = 0.0
+
+    print("-" * 100)
+    print(f"Evaluating model: {model}")
+    print(f"Using API URL: {api_url if api_url else 'No API URL'}")
+    print(f"Batch size: {batch_size}")
+    print(f"Max LiteLLM workers: {max_litellm_workers}")
+    print(f"Using prompt: {args.which_prompt}")
+    print("-" * 100)
+
+    with tqdm(total=len(testset), desc="Evaluating") as pbar:
+        for batch in batched_testset:
+            litellm_requests = []
+
+            for item in batch:
+                user_content = f"{item['input']}"
+                messages = [{"role": "system", "content": instruction}, {"role": "user", "content": user_content}]
+
+                litellm_requests.append(messages)
+
+            try:
+                responses = litellm.batch_completion(
+                    model=model,
+                    messages=litellm_requests,
+                    api_base=api_url,
+                    max_workers=max_litellm_workers,
+                    format=AnyMathsStructuredOutput.model_json_schema(),
+                    response_format={
+                        "type": "json_object",
+                        "response_schema": AnyMathsStructuredOutput.model_json_schema(),
+                        "enforce_validation": True,
+                    },
+                )
+            except litellm.exceptions.JSONSchemaValidationError as e:
+                raise e
+
+            for response, item in zip(responses, batch, strict=False):
+                correct_output_format = True
+                try:
+                    assistant_response = ast.literal_eval(response.choices[0].message.content.strip())
+                    assistant_final_answer = assistant_response["final_answer"]
+                    ground_truth = item["answer"]
+                    score = 1.0 if ground_truth in assistant_final_answer else 0.0
+                    total_score += score
+                except Exception:
+                    correct_output_format = False
+                    continue
+
+            pbar.update(len(batch))
+            pbar.set_postfix({"Score": f"{total_score} / {len(testset):.4f}"})
+
+    print("-" * 100)
+    print(f"Final score >> {total_score} / {len(testset):.4f}")
+    print("-" * 100)

mantisdk/algorithm/gepa/lib/examples/anymaths-bench/prompt-templates/instruction_prompt.txt

@@ -0,0 +1,9 @@
+You are an AI assistant that solves mathematical word problems. You will be given a question and you need to provide a step-by-step solution to the problem. Finally, you will provide the answer to the question. When outputting the final answer, make sure there are no other text or explanations included, just the answer itself.
+
+The expected output must be a JSON object with the following format:
+{
+  "final_answer": <the final answer to the question>,
+  "solution_pad": <the step-by-step solution to the problem>
+}
+
+Strictly follow the format provided above and ensure that your output is a valid JSON object. Any deviation from this format will result in an error.

mantisdk/algorithm/gepa/lib/examples/anymaths-bench/prompt-templates/optimal_prompt.txt

@@ -0,0 +1,24 @@
+You are an AI assistant that solves mathematical word problems. You will be given a question and you need to provide a step-by-step solution to the problem. Finally, you will provide the answer to the question.
+
+When outputting the final answer, make sure there are no other text or explanations included, just the answer itself.
+
+The following fields are what you need to include in your response:
+- final_answer: The final answer to the question.
+- solution_pad: The step-by-step solution to the problem.
+
+Here are specific guidelines for generating your response:
+
+1.  **Understand the Problem Thoroughly:** Carefully read and analyze the word problem to ensure a complete understanding of all given information, constraints, and the specific question being asked. Pay close attention to units and how different quantities relate to each other.
+
+2.  **Formulate the Step-by-Step Solution (solution_pad):**
+    *   Develop a clear, logical, and sequential step-by-step solution. Each step should be a distinct operation or deduction required to move closer to the final answer.
+    *   Clearly state what is being calculated or determined in each step.
+    *   Perform all necessary calculations with high precision and accuracy. Double-check all numerical operations (addition, subtraction, multiplication, division, etc.) to prevent errors.
+    *   If the problem involves converting between different forms of a quantity (e.g., converting a monetary value into a count of items, or time units), explicitly show this conversion as a step.
+        *   **Domain-Specific Interpretation Example:** If Barry has "$10.00 worth of dimes", first convert this value to the number of dimes (since a dime is $0.10, Barry has $10.00 / $0.10 = 100 dimes). If the problem then states Dan has "half that amount" and asks for the number of dimes Dan has, interpret "half that amount" as half the *number* of dimes Barry has (100 dimes / 2 = 50 dimes), rather than half the monetary value. Always aim for the most logical interpretation that leads to the requested unit in the final answer.
+    *   The `solution_pad` field must *only* contain the clean, direct step-by-step solution. Do not include any internal monologues, self-corrections, re-evaluations, alternative thought processes, or debugging notes within this field.
+
+3.  **Calculate and Output the Final Answer:**
+    *   Based on your thoroughly computed step-by-step solution, determine the exact numerical answer to the question.
+    *   The `final_answer` field must contain *only* the numerical value. Do not include any currency symbols (e.g., "$"), units (e.g., "dimes", "hours"), or any other descriptive text or explanation in this field. For example, if the answer is 4625 dollars, output `4625`. If the answer is 52 dimes, output `52`.
+    *   Ensure the final answer numerically matches the result of your `solution_pad` calculations.'

mantisdk/algorithm/gepa/lib/examples/anymaths-bench/train_anymaths.py

@@ -0,0 +1,177 @@
+def init_dataset(anymaths_dset_name: str = "openai/gsm8k"):
+    import random
+
+    from datasets import load_dataset
+
+    train_split = []
+    test_split = []
+    match anymaths_dset_name:
+        case "openai/gsm8k":
+            train_load_dataset = load_dataset(anymaths_dset_name, "main", split="train")
+            for item in train_load_dataset:
+                answer = item["answer"].split("####")[-1].strip()
+                solution = item["answer"].split("####")[0].strip()
+                question = item["question"]
+
+                train_split.append({"input": question, "additional_context": {"solution": solution}, "answer": answer})
+
+            random.Random(0).shuffle(train_split)
+
+            test_load_dataset = load_dataset(anymaths_dset_name, "main", split="test")
+            for item in test_load_dataset:
+                answer = item["answer"].split("####")[-1].strip()
+                solution = item["answer"].split("####")[0].strip()
+                question = item["question"]
+
+                test_split.append({"input": question, "answer": answer})
+
+        case "MathArena/aime_2025":
+            train_load_dataset = load_dataset("AI-MO/aimo-validation-aime", "default", split="train")
+            for item in train_load_dataset:
+                question = item["problem"]
+                solution = item["solution"]
+                answer = item["answer"]
+
+                train_split.append({"input": question, "additional_context": {"solution": solution}, "answer": answer})
+
+            random.Random(0).shuffle(train_split)
+
+            test_load_dataset = load_dataset("MathArena/aime_2025", "default", split="train")
+            for item in test_load_dataset:
+                question = item["problem"]
+                answer = item["answer"]
+
+                test_split.append({"input": question, "answer": answer})
+        case _:
+            raise ValueError(f"Unknown dataset name: {anymaths_dset_name}")
+
+    trainset = train_split[: len(train_split) // 2]
+    valset = train_split[len(train_split) // 2 :]
+    testset = test_split
+
+    return trainset, valset, testset
+
+
+if __name__ == "__main__":
+    import argparse
+    from functools import partial
+    from pathlib import Path
+
+    import litellm
+
+    from mantisdk.algorithm.gepa.lib import optimize
+    from mantisdk.algorithm.gepa.lib.adapters.anymaths_adapter import AnyMathsAdapter
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--anymaths_dset_name", type=str, default="openai/gsm8k")
+    parser.add_argument("--train_size", type=int, default=1, help="The size of the training set to use.")
+    parser.add_argument("--val_size", type=int, default=1, help="The size of the validation set to use.")
+    parser.add_argument("--test_size", type=int, default=1, help="The size of the test set to use.")
+    parser.add_argument("--base_lm", type=str, default="ollama/qwen3:4b")
+    parser.add_argument("--use_api_base", action="store_true", help="Use API base URL")
+    parser.add_argument("--api_base_url", type=str, default="http://localhost:11434")
+    parser.add_argument(
+        "--reflection_lm", type=str, default="ollama/qwen3:8b", help="The name of the reflection LM to use."
+    )
+    parser.add_argument("--use_api_reflection", action="store_true", help="Use API reflection URL")
+    parser.add_argument(
+        "--api_reflection_url",
+        type=str,
+        default="http://localhost:11434",
+        help="The API base URL for the reflection LM.",
+    )
+    parser.add_argument(
+        "--reflection_minibatch_size", type=int, default=8, help="The size of the minibatch for the reflection LM."
+    )
+    parser.add_argument("--max_litellm_workers", type=int, default=10)
+    parser.add_argument("--budget", type=int, default=500, help="The budget for the optimization process.")
+    parser.add_argument(
+        "--seed", type=int, default=0, help="The seed for the random number generator for reproducibility."
+    )
+    args = parser.parse_args()
+
+    INSTRUCTION_PROMPT_PATH = Path(__file__).parent / "prompt-templates/instruction_prompt.txt"
+
+    seed_instruction = INSTRUCTION_PROMPT_PATH.read_text()
+
+    trainset, valset, testset = init_dataset(args.anymaths_dset_name)
+
+    train_size = args.train_size
+    val_size = args.val_size
+    test_size = args.test_size
+
+    for size in map(int, [train_size, val_size, test_size]):
+        if size <= 0:
+            raise ValueError("Train, val, and test sizes must be positive integers.")
+
+    trainset = trainset[:train_size]
+    valset = valset[:val_size]
+    testset = testset[:test_size]
+
+    print("-" * 100)
+    print(f"Using dataset: {args.anymaths_dset_name}")
+    print(f"Training set size: {len(trainset)}")
+    print(f"Validation set size: {len(valset)}")
+    print(f"Test set size: {len(testset)}")
+    print("-" * 100)
+
+    base_lm = args.base_lm
+
+    reflection_lm_name = args.reflection_lm
+
+    _reflection = {"model": reflection_lm_name}
+
+    use_api_base = args.use_api_base
+    use_api_reflection = args.use_api_reflection
+
+    if use_api_base:
+        api_base = args.api_base_url
+    else:
+        api_base = None
+
+    if use_api_reflection:
+        api_reflection = args.api_reflection_url
+        _reflection["base_url"] = api_reflection
+    else:
+        api_reflection = None
+
+    _reflection_completion = partial(litellm.completion, **_reflection)
+
+    def reflection_lm(prompt: str):
+        """Call the reflection language model with the given prompt and return its content string."""
+        response = _reflection_completion(messages=[{"role": "user", "content": prompt}])
+        return response.choices[0].message.content
+
+    max_litellm_workers = args.max_litellm_workers
+    budget = args.budget
+    reflection_minibatch_size = args.reflection_minibatch_size
+    seed = args.seed
+
+    print(f"Using base LM: {base_lm}")
+    print(f"Using reflection LM: {reflection_lm_name}")
+    print(f"Using API base URL: {api_base}")
+    print(f"Using API reflection URL: {api_reflection}")
+    print(f"Reflection minibatch size: {reflection_minibatch_size}")
+    print(f"Max LiteLLM workers: {max_litellm_workers}")
+    print(f"Budget: {budget}")
+    print(f"Seed: {seed}")
+    print("-" * 100)
+
+    optimized_results = optimize(
+        seed_candidate={"instruction_prompt": seed_instruction},
+        trainset=trainset,
+        valset=valset,
+        adapter=AnyMathsAdapter(model=base_lm, api_base=api_base, max_litellm_workers=max_litellm_workers),
+        reflection_lm=reflection_lm,
+        reflection_minibatch_size=reflection_minibatch_size,
+        perfect_score=1,
+        skip_perfect_score=False,
+        use_wandb=False,
+        max_metric_calls=budget,
+        seed=seed,
+        display_progress_bar=True,
+    )
+
+    print("-" * 100)
+    print(f"Best prompt >>> {optimized_results.best_candidate}")
+    print("-" * 100)