mcpbr 0.4.15__py3-none-any.whl → 0.5.0__py3-none-any.whl

mcpbr/few_shot.py

@@ -0,0 +1,367 @@
+"""Few-shot learning support for benchmark evaluations.
+
+Provides configurable few-shot example selection with multiple strategies
+(random, similar, diverse), prompt formatting, and learning curve analysis
+to study how performance changes with varying shot counts.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class FewShotConfig:
+    """Configuration for few-shot example selection.
+
+    Attributes:
+        num_shots: Number of examples to include (0 for zero-shot).
+        selection_strategy: Strategy for selecting examples. One of
+            ``"random"``, ``"similar"``, or ``"diverse"``.
+        seed: Random seed for reproducibility. ``None`` for non-deterministic.
+    """
+
+    num_shots: int = 0
+    selection_strategy: str = "random"
+    seed: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate configuration values after initialisation."""
+        if self.num_shots < 0:
+            raise ValueError(f"num_shots must be non-negative, got {self.num_shots}")
+
+        valid_strategies = {"random", "similar", "diverse"}
+        if self.selection_strategy not in valid_strategies:
+            raise ValueError(
+                f"selection_strategy must be one of {valid_strategies}, "
+                f"got {self.selection_strategy!r}"
+            )
+
+
+def select_examples(
+    pool: list[dict[str, Any]],
+    query: dict[str, Any],
+    config: FewShotConfig,
+) -> list[dict[str, Any]]:
+    """Select few-shot examples from a pool based on the given configuration.
+
+    Args:
+        pool: List of candidate example dictionaries to select from.
+        query: The query/task dictionary that examples are being selected for.
+            Used by the ``"similar"`` strategy to find related examples.
+        config: Few-shot configuration controlling selection behaviour.
+
+    Returns:
+        List of selected example dictionaries. Length is
+        ``min(config.num_shots, len(pool))``.
+    """
+    if config.num_shots == 0 or not pool:
+        return []
+
+    num = min(config.num_shots, len(pool))
+
+    if config.selection_strategy == "random":
+        return _select_random(pool, num, config.seed)
+    elif config.selection_strategy == "similar":
+        return _select_similar(pool, query, num, config.seed)
+    elif config.selection_strategy == "diverse":
+        return _select_diverse(pool, num, config.seed)
+    else:
+        raise ValueError(f"Unknown selection strategy: {config.selection_strategy!r}")
+
+
+def format_few_shot_prompt(
+    examples: list[dict[str, Any]],
+    query: dict[str, Any],
+    template: str,
+) -> str:
+    """Format a prompt string with few-shot examples and the target query.
+
+    The *template* string should contain ``{examples}`` and ``{query}``
+    placeholders. Each example is formatted as a numbered block using the
+    example's string representation.
+
+    Args:
+        examples: List of few-shot example dictionaries.
+        query: The target query dictionary.
+        template: A format string containing ``{examples}`` and ``{query}``
+            placeholders.
+
+    Returns:
+        The fully formatted prompt string.
+    """
+    if not examples:
+        examples_text = "No examples provided."
+    else:
+        parts: list[str] = []
+        for i, example in enumerate(examples, 1):
+            parts.append(f"Example {i}:\n{_format_dict(example)}")
+        examples_text = "\n\n".join(parts)
+
+    query_text = _format_dict(query)
+
+    return template.format(examples=examples_text, query=query_text)
+
+
+def compute_learning_curve(
+    results_by_shots: dict[int, list[dict[str, Any]]],
+) -> dict[str, Any]:
+    """Analyse how performance changes with the number of few-shot examples.
+
+    For each shot count, computes accuracy (fraction of results where
+    ``"resolved"`` is truthy), average cost, and average token usage.
+
+    Args:
+        results_by_shots: Mapping from shot count (int) to a list of result
+            dictionaries. Each result dict may contain keys such as
+            ``"resolved"`` (bool), ``"cost"`` (float), and
+            ``"tokens"`` (dict with ``"input"`` and ``"output"``).
+
+    Returns:
+        Dictionary with keys:
+
+        - ``"shot_counts"``: sorted list of shot counts.
+        - ``"accuracy"``: list of accuracy values corresponding to each shot count.
+        - ``"avg_cost"``: list of average cost values.
+        - ``"avg_tokens"``: list of average total token counts.
+        - ``"num_samples"``: list of sample counts for each shot count.
+    """
+    if not results_by_shots:
+        return {
+            "shot_counts": [],
+            "accuracy": [],
+            "avg_cost": [],
+            "avg_tokens": [],
+            "num_samples": [],
+        }
+
+    sorted_shots = sorted(results_by_shots.keys())
+    accuracies: list[float] = []
+    avg_costs: list[float] = []
+    avg_tokens: list[float] = []
+    num_samples: list[int] = []
+
+    for shot_count in sorted_shots:
+        results = results_by_shots[shot_count]
+        n = len(results)
+        num_samples.append(n)
+
+        if n == 0:
+            accuracies.append(0.0)
+            avg_costs.append(0.0)
+            avg_tokens.append(0.0)
+            continue
+
+        # Accuracy
+        resolved = sum(1 for r in results if r.get("resolved"))
+        accuracies.append(resolved / n)
+
+        # Average cost
+        total_cost = sum(r.get("cost", 0.0) for r in results)
+        avg_costs.append(total_cost / n)
+
+        # Average tokens
+        total_tokens = 0
+        for r in results:
+            tokens = r.get("tokens", {})
+            total_tokens += tokens.get("input", 0) + tokens.get("output", 0)
+        avg_tokens.append(total_tokens / n)
+
+    return {
+        "shot_counts": sorted_shots,
+        "accuracy": accuracies,
+        "avg_cost": avg_costs,
+        "avg_tokens": avg_tokens,
+        "num_samples": num_samples,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _format_dict(d: dict[str, Any]) -> str:
+    """Format a dictionary as a human-readable key-value block.
+
+    Args:
+        d: Dictionary to format.
+
+    Returns:
+        Multi-line string with one ``key: value`` pair per line.
+    """
+    lines = [f"  {k}: {v}" for k, v in d.items()]
+    return "\n".join(lines) if lines else "  (empty)"
+
+
+def _select_random(
+    pool: list[dict[str, Any]],
+    num: int,
+    seed: int | None,
+) -> list[dict[str, Any]]:
+    """Randomly select *num* examples from *pool*.
+
+    Args:
+        pool: Candidate examples.
+        num: Number of examples to select.
+        seed: Random seed for reproducibility.
+
+    Returns:
+        Randomly selected examples.
+    """
+    rng = random.Random(seed)
+    return rng.sample(pool, num)
+
+
+def _select_similar(
+    pool: list[dict[str, Any]],
+    query: dict[str, Any],
+    num: int,
+    seed: int | None,
+) -> list[dict[str, Any]]:
+    """Select examples most similar to *query* by shared category/tags.
+
+    Similarity is computed as the number of matching metadata values
+    for the keys ``"category"``, ``"tags"``, ``"language"``, and ``"difficulty"``.
+    For the ``"tags"`` key the score is the size of the intersection.
+
+    When there are ties in similarity score, a seeded RNG is used to
+    shuffle tied candidates for deterministic tie-breaking.
+
+    Args:
+        pool: Candidate examples.
+        query: The query dictionary to compare against.
+        num: Number of examples to select.
+        seed: Random seed for reproducibility in tie-breaking.
+
+    Returns:
+        Examples sorted by descending similarity, with ties broken
+        deterministically when *seed* is provided.
+    """
+    rng = random.Random(seed)
+
+    scored: list[tuple[float, int, dict[str, Any]]] = []
+    for idx, example in enumerate(pool):
+        score = _similarity_score(example, query)
+        scored.append((score, idx, example))
+
+    # Group by score and shuffle within each group for deterministic tie-breaking
+    # Sort descending by score first
+    scored.sort(key=lambda x: x[0], reverse=True)
+
+    # Build groups of tied scores
+    groups: list[list[tuple[float, int, dict[str, Any]]]] = []
+    current_group: list[tuple[float, int, dict[str, Any]]] = []
+    current_score: float | None = None
+
+    for item in scored:
+        if current_score is None or item[0] == current_score:
+            current_group.append(item)
+            current_score = item[0]
+        else:
+            groups.append(current_group)
+            current_group = [item]
+            current_score = item[0]
+    if current_group:
+        groups.append(current_group)
+
+    # Shuffle within each group and flatten
+    result: list[dict[str, Any]] = []
+    for group in groups:
+        rng.shuffle(group)
+        for _, _, example in group:
+            result.append(example)
+            if len(result) == num:
+                return result
+
+    return result[:num]
+
+
+def _similarity_score(example: dict[str, Any], query: dict[str, Any]) -> float:
+    """Compute a simple similarity score between an example and a query.
+
+    Checks the following keys:
+    - ``"category"``: +1 if equal.
+    - ``"language"``: +1 if equal.
+    - ``"difficulty"``: +1 if equal.
+    - ``"tags"``: +1 for each shared tag (set intersection size).
+
+    Args:
+        example: Candidate example dictionary.
+        query: Target query dictionary.
+
+    Returns:
+        Non-negative similarity score.
+    """
+    score = 0.0
+
+    for key in ("category", "language", "difficulty"):
+        if key in example and key in query and example[key] == query[key]:
+            score += 1.0
+
+    example_tags = set(example.get("tags", []))
+    query_tags = set(query.get("tags", []))
+    score += len(example_tags & query_tags)
+
+    return score
+
+
+def _select_diverse(
+    pool: list[dict[str, Any]],
+    num: int,
+    seed: int | None,
+) -> list[dict[str, Any]]:
+    """Select diverse examples covering different categories.
+
+    Groups candidates by their ``"category"`` key (falling back to
+    ``"_uncategorized_"``), then round-robins across categories to
+    select *num* examples. Within each category, examples are shuffled
+    using the provided *seed* for reproducibility.
+
+    Args:
+        pool: Candidate examples.
+        num: Number of examples to select.
+        seed: Random seed for reproducibility.
+
+    Returns:
+        Diverse selection of examples from different categories.
+    """
+    rng = random.Random(seed)
+
+    # Group by category
+    categories: dict[str, list[dict[str, Any]]] = {}
+    for example in pool:
+        cat = example.get("category", "_uncategorized_")
+        if cat not in categories:
+            categories[cat] = []
+        categories[cat].append(example)
+
+    # Shuffle within each category for fairness
+    for cat_examples in categories.values():
+        rng.shuffle(cat_examples)
+
+    # Sort category keys for deterministic ordering
+    sorted_cats = sorted(categories.keys())
+
+    # Track current index within each category's shuffled list
+    cat_indices: dict[str, int] = {cat: 0 for cat in sorted_cats}
+
+    result: list[dict[str, Any]] = []
+    while len(result) < num:
+        added_any = False
+        for cat in sorted_cats:
+            if len(result) >= num:
+                break
+            idx = cat_indices[cat]
+            if idx < len(categories[cat]):
+                result.append(categories[cat][idx])
+                cat_indices[cat] = idx + 1
+                added_any = True
+
+        # If no category had remaining items, we've exhausted the pool
+        if not added_any:
+            break
+
+    return result