thinkbooster 0.1.0__py3-none-any.whl

llm_tts/datasets/__init__.py

@@ -0,0 +1,46 @@
+"""Dataset loaders for various benchmarks."""
+
+from .gsm8k import (
+    evaluate_gsm8k_answer,
+    extract_answer_from_gsm8k,
+    format_gsm8k_for_deepconf,
+    load_gsm8k,
+)
+from .human_eval_plus import create_evalplus_samples as create_human_eval_plus_samples
+from .human_eval_plus import (
+    extract_code_from_response as extract_code_from_response_human_eval,
+)
+from .human_eval_plus import (
+    format_human_eval_prompt,
+)
+from .human_eval_plus import load_evalplus_samples as load_human_eval_plus_samples
+from .human_eval_plus import (
+    load_human_eval_plus,
+)
+from .mbpp_plus import (
+    create_evalplus_samples,
+    extract_code_from_response,
+    format_mbpp_prompt,
+    load_evalplus_samples,
+    load_mbpp_plus,
+)
+
+__all__ = [
+    # GSM8K
+    "load_gsm8k",
+    "evaluate_gsm8k_answer",
+    "extract_answer_from_gsm8k",
+    "format_gsm8k_for_deepconf",
+    # MBPP+
+    "load_mbpp_plus",
+    "extract_code_from_response",
+    "format_mbpp_prompt",
+    "create_evalplus_samples",
+    "load_evalplus_samples",
+    # HumanEval+
+    "load_human_eval_plus",
+    "extract_code_from_response_human_eval",
+    "format_human_eval_prompt",
+    "create_human_eval_plus_samples",
+    "load_human_eval_plus_samples",
+]

llm_tts/datasets/gsm8k.py

@@ -0,0 +1,168 @@
+"""
+GSM8K dataset loader and preprocessing for DeepConf evaluation.
+
+GSM8K (Grade School Math 8K) is a dataset of 8.5K grade school math word problems.
+Each problem requires multi-step reasoning to solve.
+"""
+
+import logging
+from typing import Dict, List, Optional
+
+from datasets import load_dataset
+
+log = logging.getLogger(__name__)
+
+
+def extract_answer_from_gsm8k(solution: str) -> str:
+    """
+    Extract the final numerical answer from GSM8K solution format.
+
+    GSM8K solutions end with "#### {answer}" format.
+
+    Args:
+        solution: The solution string from GSM8K
+
+    Returns:
+        The numerical answer as a string
+    """
+    if "####" in solution:
+        answer = solution.split("####")[-1].strip()
+        # Remove commas from numbers (e.g., "1,000" -> "1000")
+        answer = answer.replace(",", "")
+        return answer
+    return ""
+
+
+def format_gsm8k_for_deepconf(question: str, answer: str) -> Dict[str, str]:
+    """
+    Format GSM8K data for DeepConf evaluation.
+
+    Converts GSM8K format to the format expected by DeepConf:
+    - Question stays as is
+    - Answer is extracted from "#### X" format
+    - Expected output format is \\boxed{X}
+
+    Args:
+        question: The question text
+        answer: The solution text (includes #### {answer})
+
+    Returns:
+        Dict with 'question' and 'answer' keys
+    """
+    extracted_answer = extract_answer_from_gsm8k(answer)
+
+    return {
+        "question": question.strip(),
+        "answer": extracted_answer,
+        "original_solution": answer,  # Keep for reference
+    }
+
+
+def load_gsm8k(
+    split: str = "test",
+    subset_size: Optional[int] = None,
+    cache_dir: Optional[str] = None,
+) -> List[Dict[str, str]]:
+    """
+    Load GSM8K dataset and format for DeepConf.
+
+    Args:
+        split: Dataset split ('train' or 'test')
+        subset_size: If provided, only load first N examples
+        cache_dir: Cache directory for HuggingFace datasets
+
+    Returns:
+        List of dicts with 'question' and 'answer' keys
+    """
+    log.info(f"Loading GSM8K dataset (split={split})...")
+
+    # Load from HuggingFace
+    dataset = load_dataset("openai/gsm8k", "main", split=split, cache_dir=cache_dir)
+
+    # Take subset if requested
+    if subset_size is not None:
+        dataset = dataset.select(range(min(subset_size, len(dataset))))
+        log.info(f"Using subset of {len(dataset)} examples")
+
+    # Format for DeepConf
+    formatted_data = []
+    for item in dataset:
+        formatted = format_gsm8k_for_deepconf(
+            question=item["question"], answer=item["answer"]
+        )
+        formatted_data.append(formatted)
+
+    log.info(f"Loaded {len(formatted_data)} GSM8K examples")
+
+    return formatted_data
+
+
+def evaluate_gsm8k_answer(predicted: str, ground_truth: str) -> bool:
+    """
+    Evaluate if predicted answer matches ground truth for GSM8K.
+
+    Handles:
+    - Numeric comparison (with tolerance for floats)
+    - String normalization (strip whitespace, lowercase)
+    - Comma removal from numbers
+
+    Args:
+        predicted: Predicted answer (extracted from \\boxed{})
+        ground_truth: Ground truth answer
+
+    Returns:
+        True if answers match, False otherwise
+    """
+    # Normalize both answers
+    pred_clean = predicted.strip().replace(",", "").lower()
+    gt_clean = ground_truth.strip().replace(",", "").lower()
+
+    # Direct string match
+    if pred_clean == gt_clean:
+        return True
+
+    # Try numeric comparison
+    try:
+        pred_num = float(pred_clean)
+        gt_num = float(gt_clean)
+
+        # Use relative tolerance for floating point comparison
+        return abs(pred_num - gt_num) < 1e-6 * max(abs(pred_num), abs(gt_num), 1)
+    except (ValueError, TypeError):
+        pass
+
+    return False
+
+
+if __name__ == "__main__":
+    # Test loading
+    logging.basicConfig(level=logging.INFO)
+
+    print("\n=== Testing GSM8K loader ===\n")
+
+    # Load small subset
+    data = load_gsm8k(split="test", subset_size=5)
+
+    print(f"Loaded {len(data)} examples\n")
+
+    for i, item in enumerate(data[:3]):
+        print(f"Example {i+1}:")
+        print(f"  Question: {item['question'][:100]}...")
+        print(f"  Answer: {item['answer']}")
+        print(f"  Original: {item['original_solution'][:80]}...")
+        print()
+
+    # Test answer evaluation
+    print("\n=== Testing answer evaluation ===\n")
+    test_cases = [
+        ("70", "70", True),
+        ("70", "70.0", True),
+        ("1000", "1,000", True),
+        ("70", "71", False),
+        ("abc", "ABC", True),
+    ]
+
+    for pred, gt, expected in test_cases:
+        result = evaluate_gsm8k_answer(pred, gt)
+        status = "✓" if result == expected else "✗"
+        print(f"{status} '{pred}' vs '{gt}': {result} (expected {expected})")

llm_tts/datasets/human_eval_plus.py

@@ -0,0 +1,266 @@
+"""
+HumanEval+ dataset loader and utilities.
+
+HumanEval+ is an enhanced version of HumanEval with 80x more test cases
+for rigorous evaluation of code generation.
+
+Dataset: https://huggingface.co/datasets/evalplus/humanevalplus
+EvalPlus: https://github.com/evalplus/evalplus
+"""
+
+import json
+import logging
+import re
+from typing import Any, Dict, List, Optional
+
+from evalplus.data import get_human_eval_plus, write_jsonl
+
+log = logging.getLogger(__name__)
+
+
+def load_human_eval_plus(
+    subset_size: Optional[int] = None,
+) -> List[Dict[str, Any]]:
+    """
+    Load HumanEval+ dataset using evalplus API.
+
+    Args:
+        subset_size: If provided, only load first N examples
+
+    Returns:
+        List of dicts with formatted data for the evaluation pipeline
+    """
+    return _load_from_evalplus(subset_size)
+
+
+def _load_from_evalplus(subset_size: Optional[int] = None) -> List[Dict[str, Any]]:
+    """Load HumanEval+ using evalplus API.
+
+    Formats prompts to match EvalPlus official methodology:
+    - instruction_prefix + code block with docstring
+    """
+    log.info("Loading HumanEval+ using evalplus API...")
+
+    # EvalPlus instruction prefix for chat/instruction models
+    INSTRUCTION_PREFIX = (
+        "Please provide a self-contained Python script that solves the "
+        "following problem in a markdown code block:"
+    )
+
+    problems = get_human_eval_plus()
+    formatted_data = []
+
+    for task_id, problem in problems.items():
+        # Format prompt exactly like EvalPlus does for chat models:
+        # instruction_prefix + "\n```python\n" + prompt + "\n```"
+        raw_prompt = problem["prompt"].strip()
+        formatted_prompt = f"{INSTRUCTION_PREFIX}\n```python\n{raw_prompt}\n```"
+
+        formatted = {
+            # Standard fields for the evaluation pipeline
+            "question": formatted_prompt,
+            "answer": problem["canonical_solution"],
+            # HumanEval+ specific fields
+            "task_id": task_id,
+            "entry_point": problem.get(
+                "entry_point", _extract_function_name(raw_prompt)
+            ),
+            "prompt": raw_prompt,  # Original prompt (function signature + docstring)
+            "base_input": problem.get("base_input", []),
+            "plus_input": problem.get("plus_input", []),
+            "atol": problem.get("atol", 0),
+            "contract": problem.get("contract", ""),
+        }
+        formatted_data.append(formatted)
+
+        if subset_size and len(formatted_data) >= subset_size:
+            break
+
+    log.info(f"Loaded {len(formatted_data)} HumanEval+ problems via evalplus API")
+    return formatted_data
+
+
+def _extract_function_name(prompt: str) -> str:
+    """Extract function name from HumanEval prompt."""
+    # HumanEval prompts typically start with function signature
+    # Pattern: "def function_name("
+    match = re.search(r"def (\w+)\s*\(", prompt)
+    if match:
+        return match.group(1)
+
+    # Default fallback
+    return "solution"
+
+
+def extract_code_from_response(response: str) -> str:
+    """
+    Extract Python code from model response.
+
+    Handles various formats:
+    - Code blocks with ```python or ``` markers
+    - Raw code
+    - Code with explanation
+
+    Args:
+        response: Model's response text
+
+    Returns:
+        Extracted code string
+    """
+    # Try to extract from code blocks first
+    # Match ```python ... ``` or ``` ... ``` blocks
+    code_block_pattern = r"```(?:python)?\s*\n(.*?)```"
+    code_blocks = re.findall(code_block_pattern, response, re.DOTALL)
+
+    if code_blocks:
+        # Return the last code block (usually the final solution)
+        code = code_blocks[-1].strip()
+        # Handle malformed code blocks where "python" appears on its own line
+        # Some models output "```\npython\n..." instead of "```python\n..."
+        if code.startswith("python\n"):
+            code = code[7:]  # Remove "python\n"
+        elif code.startswith("python3\n"):
+            code = code[8:]  # Remove "python3\n"
+        elif code.startswith("Python\n"):
+            code = code[7:]  # Remove "Python\n"
+        return code.strip()
+
+    # Try to find function definition
+    # Look for def ... up to the next blank line or end
+    func_pattern = r"(def \w+\s*\([^)]*\):.*?)(?:\n\n|\Z)"
+    func_matches = re.findall(func_pattern, response, re.DOTALL)
+
+    if func_matches:
+        return func_matches[-1].strip()
+
+    # Return the response as-is (might be raw code)
+    return response.strip()
+
+
+def format_human_eval_prompt(
+    problem: Dict[str, Any],
+    prompt_template: Optional[str] = None,
+) -> str:
+    """
+    Format a HumanEval+ problem into a prompt for the model.
+
+    Args:
+        problem: A formatted HumanEval+ problem dict
+        prompt_template: Optional template with {prompt}, {entry_point} placeholders
+
+    Returns:
+        Formatted prompt string
+    """
+    if prompt_template:
+        return prompt_template.format(
+            prompt=problem["question"],
+            entry_point=problem.get("entry_point", "solution"),
+            task_id=problem.get("task_id", ""),
+        )
+
+    # Default formatting - just return the prompt
+    return problem["question"]
+
+
+def create_evalplus_samples(
+    results: List[Dict[str, Any]],
+    output_path: str,
+) -> None:
+    """
+    Create a samples file in EvalPlus format for evaluation.
+
+    The format expected by evalplus is JSONL with:
+    - task_id: str
+    - solution: str (the complete solution code)
+
+    Args:
+        results: List of result dicts with 'task_id' and generated code
+        output_path: Path to save the samples file
+    """
+    log.info(f"Saving {len(results)} samples to {output_path}")
+
+    samples = []
+    for result in results:
+        task_id = result.get("task_id", "")
+        # Get the generated code from various possible fields
+        solution = (
+            result.get("generated_code")
+            or result.get("extracted_answer")
+            or result.get("generated_answer")
+            or result.get("generated_trajectory", "")
+        )
+
+        # Extract code if it contains markdown
+        if "```" in solution:
+            solution = extract_code_from_response(solution)
+
+        samples.append(
+            {
+                "task_id": task_id,
+                "solution": solution,
+            }
+        )
+
+    write_jsonl(output_path, samples)
+
+    log.info(f"Samples saved to {output_path}")
+
+
+def load_evalplus_samples(path: str) -> List[Dict[str, Any]]:
+    """
+    Load samples from an EvalPlus format file.
+
+    Args:
+        path: Path to the samples JSONL file
+
+    Returns:
+        List of sample dictionaries
+    """
+    samples = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                samples.append(json.loads(line))
+    return samples
+
+
+if __name__ == "__main__":
+    # Test loading
+    logging.basicConfig(level=logging.INFO)
+
+    print("\n=== Testing HumanEval+ loader ===\n")
+
+    # Load small subset
+    data = load_human_eval_plus(subset_size=5)
+
+    print(f"Loaded {len(data)} problems\n")
+
+    for i, item in enumerate(data[:3]):
+        print(f"Problem {i + 1}:")
+        print(f"  Task ID: {item['task_id']}")
+        print(f"  Entry point: {item['entry_point']}")
+        print(f"  Prompt: {item['question'][:100]}...")
+        print(f"  Solution preview: {item['answer'][:80]}...")
+        print()
+
+    # Test code extraction
+    print("\n=== Testing code extraction ===\n")
+
+    test_response = """
+Here's the solution:
+
+```python
+def has_close_elements(numbers: List[float], threshold: float) -> bool:
+    for i in range(len(numbers)):
+        for j in range(i + 1, len(numbers)):
+            if abs(numbers[i] - numbers[j]) < threshold:
+                return True
+    return False
+```
+
+This function checks if any two elements are closer than the threshold.
+"""
+
+    extracted = extract_code_from_response(test_response)
+    print(f"Extracted code:\n{extracted}")