adaptive-harmony 0.1.23__py3-none-any.whl

adaptive_harmony/core/utils.py

@@ -0,0 +1,365 @@
+import asyncio
+import functools
+import hashlib
+import itertools
+import json
+import random
+from pathlib import Path
+from typing import Any, Callable, Coroutine, Dict, Iterator, List, NamedTuple, Sequence, TypedDict, TypeVar
+
+import numpy as np
+from loguru import logger
+
+from adaptive_harmony import InferenceModel, StringThread, TrainingModel
+from adaptive_harmony.core.rich_counter import ProgressCounter, get_progress_counter_or_wrapper
+from adaptive_harmony.metric_logger import Logger, StdoutLogger
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+
+async def wrap_coroutine_with_progress[T](coroutine: Coroutine[Any, Any, T], progress_counter: ProgressCounter) -> T:
+    try:
+        return await coroutine
+    finally:
+        progress_counter.increment_total_counter()
+
+
+async def async_map_batch[S, T](
+    f: Callable[[S], Coroutine[Any, Any, T]],
+    data: Iterator[S],
+    batch_size: int,
+    max_failure_fraction: float = 0.5,
+) -> List[T]:
+    """
+    Process items from an iterator in batches using concurrent coroutines.
+
+    This function processes items from an iterator in batches, executing the
+    provided coroutine function concurrently for each item. It excludes failing
+    samples until it can create a new batch of results of size # batch size.
+    If more than max_failure_fraction % of # batch size tasks fail in the process
+    of creating a new batch, the function will raise the last exception encountered.
+    Results are not ordered.
+
+    Args:
+        f: Coroutine function to apply to each item
+        data: Iterator of items to process
+        batch_size: Number of items to process in each batch
+
+    Returns:
+        List of results from successful task executions
+
+    Note:
+        - Failed tasks are not retried
+        - If more than max_failure_fraction of # batch size tasks fail, the function fails
+        - Tasks are automatically cancelled if the function exits early
+    """
+    batch_items_from_iterator = list(itertools.islice(data, batch_size))
+    num_items = len(batch_items_from_iterator)
+
+    async with get_progress_counter_or_wrapper(f"async_map_batch({f.__name__})", batch_size) as counter:
+        final_results: list[Any] = [None] * num_items
+        active_tasks_this_batch: Dict[asyncio.Task, int] = {}
+
+        num_retries = 0
+
+        for i, item_value in enumerate(batch_items_from_iterator):
+            task: asyncio.Task[T] = asyncio.create_task(wrap_coroutine_with_progress(f(item_value), counter))
+            counter.register_task(task)
+            active_tasks_this_batch[task] = i
+
+        try:
+            while active_tasks_this_batch:
+                done_tasks, _ = await asyncio.wait(active_tasks_this_batch.keys(), return_when=asyncio.FIRST_COMPLETED)
+
+                for task_item in done_tasks:
+                    original_batch_slot_idx = active_tasks_this_batch.pop(task_item)
+
+                    try:
+                        result: T = await task_item
+                        final_results[original_batch_slot_idx] = result
+                    except Exception as ex:
+                        try:
+                            if num_retries > batch_size * max_failure_fraction:
+                                # if more than 50% of a batch fail we'll just go on.
+                                raise ex
+
+                            logger.debug(ex)
+                            retry_item_value: S = next(data)
+                            new_retry_task: asyncio.Task[T] = asyncio.create_task(
+                                wrap_coroutine_with_progress(f(retry_item_value), counter)
+                            )
+                            active_tasks_this_batch[new_retry_task] = original_batch_slot_idx
+                            num_retries += 1
+                        except StopIteration:
+                            ...
+        finally:
+            tasks_to_cancel = list(active_tasks_this_batch.keys())
+            for task_to_cancel in tasks_to_cancel:
+                task_to_cancel.cancel()
+
+            if tasks_to_cancel:
+                await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
+
+        if num_retries > 0:
+            print(f"WARNING: had to retry {num_retries} times to get a batch of {batch_size}")
+        ret = [res for res in final_results if res is not None]
+
+        print(f"Final number tasks with non-None results: {len(ret)}")
+
+        return ret
+
+
+def hash_hyperparams(include: set[str]):
+    """
+    A decorator that computes a hash of specified hyperparameters and stores it on `self._hyperparams_hash`.
+
+    Must be used on an `__init__` method. Only parameters listed in `include` will be hashed.
+    Non-serializable values are converted to their string representation.
+
+    Args:
+        include: Set of parameter names to include in the hash.
+
+    Example:
+        @hash_hyperparams(include={"lr", "batch_size", "epochs"})
+        def __init__(self, lr, batch_size, epochs, logger, ...):
+            ...
+    """
+    import inspect
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            sig = inspect.signature(func)
+            bound_args = sig.bind_partial(self, *args, **kwargs)
+            bound_args.apply_defaults()
+            all_args = bound_args.arguments
+
+            hyperparams = {}
+            for key in include:
+                if key in all_args:
+                    value = all_args[key]
+                    try:
+                        json.dumps(value)
+                        hyperparams[key] = value
+                    except (TypeError, OverflowError):
+                        hyperparams[key] = repr(value)
+
+            serialized = json.dumps(hyperparams, sort_keys=True)
+            self._hyperparams_hash = hashlib.sha256(serialized.encode()).hexdigest()[:16]
+
+            return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def log_args(func):
+    """
+    A Python decorator that logs the arguments of the decorated function
+    to experiment tracking tools (wandb, mlflow) or stdout.
+
+    Attempts to log to wandb if available and initialized, then to mlflow
+    if available and has an active run. If neither is available, logs to stdout.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        import inspect
+
+        # Helper to check serializability and prepare value
+        def prepare_value(value):
+            # we need to log the model builder args here because they are not serializable by default
+            if isinstance(value, list) and len(value) > 100:
+                # exclude long lists since we want to skip datasets
+                return None
+            if isinstance(value, InferenceModel) or isinstance(value, TrainingModel):
+                return value.get_builder_args()  # type: ignore PyRight being dumb
+            else:
+                # Check if the value itself is a complex object that might not be fully serializable
+                try:
+                    json.dumps({"test_key": value})
+                    return value
+                except (TypeError, OverflowError):
+                    return None
+
+        # Get function arguments once
+        sig = inspect.signature(func)
+        bound_args = sig.bind_partial(*args, **kwargs)
+        bound_args.apply_defaults()
+        all_args = bound_args.arguments
+
+        # find the loggers that are given to recipe, if None are found, we will log to stdout
+        loggers = [v for v in all_args.values() if isinstance(v, Logger)]
+        if not loggers:
+            loggers.append(StdoutLogger())
+
+        # get loggable args only
+        loggable_args = {k: new_v for k, v in all_args.items() if (new_v := prepare_value(v)) is not None}
+
+        for logger_instance in loggers:
+            logger_instance.log_config(loggable_args)
+
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+async def async_map[S, T](f: Callable[[S], Coroutine[Any, Any, T]], data: Sequence[S]) -> list[T]:
+    """
+    Process all items in an iterable concurrently using the provided coroutine function.
+
+    Args:
+        f: Coroutine function to apply to each item
+        data: Iterable of items to process
+
+    Returns:
+        List of results from all task executions
+    """
+
+    # Check if a Progress bar is already active
+    async with get_progress_counter_or_wrapper(f"async_map({f.__name__})", len(list(data))) as counter:
+        all_tasks = [asyncio.create_task(wrap_coroutine_with_progress(f(item), counter)) for item in data]
+        for t in all_tasks:
+            counter.register_task(t)
+        results = await asyncio.gather(*all_tasks)
+    return results
+
+
+async def async_map_fallible[S, T](f: Callable[[S], Coroutine[Any, Any, T]], data: Sequence[S]) -> list[T]:
+    """
+    Process all items in an iterable concurrently using the provided coroutine function.
+
+    Args:
+        f: Coroutine function to apply to each item
+        data: Iterable of items to process
+
+    Returns:
+        List of results from all task executions
+    """
+
+    async def wrap_coroutine_with_error_handling(coro: Coroutine[Any, Any, T]) -> tuple[T, bool]:
+        try:
+            result = await coro
+            return result, True
+        except Exception:
+            return None, False  # type: ignore
+
+    async with get_progress_counter_or_wrapper(f"async_map_fallible({f.__name__})", len(list(data))) as counter:
+        all_tasks = [
+            asyncio.create_task(wrap_coroutine_with_error_handling(wrap_coroutine_with_progress(f(item), counter)))
+            for item in data
+        ]
+        for t in all_tasks:
+            counter.register_task(t)
+        results = await asyncio.gather(*all_tasks)
+
+    return [result for result, success in results if success]
+
+
+def get_minibatches[T](dataset: list[T], mini_batch_size: int, number_of_epochs: int) -> list[list[T]]:
+    all_batches: list[list[T]] = []
+
+    for _ in range(number_of_epochs):
+        shuffled_dataset = random.sample(dataset, k=len(dataset))
+
+        epoch_batches: list[list[T]] = []
+        for i in range(0, len(shuffled_dataset), mini_batch_size):
+            batch = shuffled_dataset[i : i + mini_batch_size]
+            epoch_batches.append(batch)
+        all_batches.extend(epoch_batches)
+
+    return all_batches
+
+
+def sample_data[T](data: list[T], epochs: float) -> list[T]:
+    num_samples = len(data) * epochs
+    return [data[x] for x in np.random.permutation(len(data))[: int(num_samples)]]
+
+
+def weighted_mean(values: list[list[float]], weights: list[list[float]]) -> float:
+    return np.average(np.concatenate(values), weights=np.concatenate(weights)).item()
+
+
+def stringify_thread(thread: StringThread, sep: str = "\n\n") -> str:
+    """Convert StringThread to readable text format."""
+    turns = thread.get_turns()
+    return sep.join([f"[{turn.role}]\n{turn.content}" for turn in turns])
+
+
+class SingleTurnShot(TypedDict):
+    user: dict[str, str]
+    assistant: dict[str, str]
+
+
+class TurnTemplates(NamedTuple):
+    system: str | None
+    user: str | None
+    assistant: str | None
+    shots: list[SingleTurnShot] | None
+
+
+def turn_templates_from_dir(root_dir: str) -> TurnTemplates:
+    """
+    Returns system, user and assistant turn string templates from a directory, as well as a list of shot dicts.
+    Expects files to be named system.md, user.md, assistant.md and shots.jsonl.
+    Returns None for any turn template file that does not exist.
+    """
+    root_path = Path(root_dir)
+    expected_files = ["system.md", "user.md", "assistant.md", "shots.jsonl"]
+    missing_templates = []
+    turn_templates: list[str | list[SingleTurnShot] | None] = []
+
+    for file in expected_files:
+        path = root_path / file
+        if not path.exists():
+            missing_templates.append(file)
+            turn_templates.append(None)
+        else:
+            if file == "shots.jsonl":
+                shots = []
+                for line in path.read_text().splitlines():
+                    data = json.loads(line)
+                    shot = SingleTurnShot(user=data["user"], assistant=data["assistant"])
+                    shots.append(shot)
+                turn_templates.append(shots)
+            else:
+                turn_templates.append(path.read_text())
+
+    # Ensure proper typing: first 3 are str|None, last is list[SingleTurnShot]|None
+    system, user, assistant, shots = turn_templates
+    return TurnTemplates(
+        system=system if isinstance(system, str) else None,
+        user=user if isinstance(user, str) else None,
+        assistant=assistant if isinstance(assistant, str) else None,
+        shots=shots if isinstance(shots, list) else None,
+    )
+
+
+def hash_dataset(dataset: Sequence[StringThread], num_samples: int = 5) -> str:
+    """Compute a hash of dataset for quick unsafe comparison.
+
+    Hashes: dataset length + first N elements + last N elements
+    This catches most dataset changes without processing the entire dataset.
+
+    Args:
+        dataset: List of dataset items
+        num_samples: Number of elements to sample from start/end (default 5)
+
+    Returns:
+        SHA256 hash
+    """
+    hasher = hashlib.sha256()
+
+    hasher.update(str(len(dataset)).encode())
+
+    num_samples = min(num_samples, len(dataset) // 2)
+
+    for item in dataset[:num_samples]:
+        hasher.update(str(item).encode())
+
+    for item in dataset[-num_samples:]:
+        hasher.update(str(item).encode())
+
+    return hasher.hexdigest()

adaptive_harmony/environment/__init__.py

@@ -0,0 +1,8 @@
+from .environment import Environment, EnvironmentFactory, TrajectoryScore, TurnScore
+
+__all__ = [
+    "Environment",
+    "EnvironmentFactory",
+    "TrajectoryScore",
+    "TurnScore",
+]

adaptive_harmony/environment/environment.py

@@ -0,0 +1,121 @@
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from dataclasses import dataclass
+from numbers import Number
+from typing import Any, Callable
+
+import numpy as np
+
+from adaptive_harmony import StringThread
+from adaptive_harmony.logging_table import Table
+
+
+@dataclass
+class TurnScore:
+    score: float
+    num_assistant_turns: int
+
+
+@dataclass
+class TrajectoryScore:
+    scores: list[TurnScore]
+
+    def to_turn_scores(self) -> list[float]:
+        return [turn_score.score for turn_score in self.scores for _ in range(turn_score.num_assistant_turns)]
+
+    @property
+    def cumulative_score(self) -> float:
+        return sum([turn_score.score for turn_score in self.scores])
+
+
+class Environment(ABC):
+    """
+    Environment to inherit from when building an environment.
+    """
+
+    def __init__(self, logging_function: Callable):
+        self.logging_function = logging_function
+
+    @abstractmethod
+    async def react_to(self, thread: StringThread) -> list[tuple[str, str]] | TrajectoryScore:
+        """Returns either [("tool", tool response), ...] or [("user", user question)] or TrajectoryScore when DONE."""
+        pass
+
+    async def bootstrap_prompt(self, thread: StringThread) -> StringThread:
+        return thread
+
+    async def generate_trajectory_and_grade(
+        self, model, initial_thread: StringThread
+    ) -> tuple[StringThread, TrajectoryScore]:
+        """Generate a full trajectory by interacting with the model until termination."""
+        thread = await self.bootstrap_prompt(initial_thread)
+
+        while True:
+            # Generate model response
+            thread = await model.generate(thread)
+
+            # Get environment reaction
+            env_response = await self.react_to(thread)
+
+            # If we got a score, we're done
+            if isinstance(env_response, TrajectoryScore):
+                return thread, env_response
+
+            # Otherwise, add the environment responses to the thread
+            for role, content in env_response:
+                if role == "tool":
+                    thread = thread.tool(content)
+                elif role == "user":
+                    thread = thread.user(content)
+                else:
+                    raise ValueError(f"Unknown role: {role}")
+
+
+class EnvironmentFactory(ABC):
+    """
+    Abstract class to build environments. It is necessary because each trajectory must have its own unique Environment
+    """
+
+    def __init__(self, logging_name: str | None = None):
+        self._logs: dict[str, list] = defaultdict(list)
+        self.logging_name = logging_name
+
+    def add_log(self, key, new_value) -> None:
+        """Add a log entry to the scorer's log collection."""
+        self._logs[key].append(new_value)
+
+    def get_logs(self, clear: bool = False) -> dict[str, float | Table]:
+        """
+        Get aggregated logs from all score calls.
+        Base implementation computes statistics for "score" keys in individual logs.
+        If there are none, returns empty dict.
+        """
+        if not self._logs:
+            return {}
+
+        logs = {}
+        for k, v in self._logs.items():
+            if isinstance(v[0], Number):
+                logs[k] = np.asarray(v).mean()
+            elif isinstance(v[0], Table):
+                headers = v[0].headers
+                assert all(table.headers == headers for table in v)
+                overall_table = Table(headers)
+                for table in v:
+                    overall_table.add_rows(table.rows)
+                logs[k] = overall_table
+            else:
+                raise ValueError(f"Unknown type: {type(v[0])}")
+
+        if clear:
+            self.clear_logs()
+        return logs
+
+    def clear_logs(self) -> None:
+        """
+        Clear all accumulated logs.
+        """
+        self._logs.clear()
+
+    @abstractmethod
+    def create_environment(self, metadata: Any) -> Environment: ...

adaptive_harmony/evaluation/__init__.py

@@ -0,0 +1 @@
+from .evaluation_artifact import EvaluationArtifact as EvaluationArtifact

adaptive_harmony/evaluation/evaluation_artifact.py

@@ -0,0 +1,67 @@
+import logging
+import uuid
+from typing import List, Self
+
+from harmony_client import (
+    EvalSample,
+    EvaluationArtifactBase,
+)
+from harmony_client.runtime.context import RecipeContext
+
+logger = logging.getLogger(__name__)
+
+
+class EvaluationArtifact:
+    def __init__(self, name: str, ctx: RecipeContext) -> None:
+        artifact_id = str(uuid.uuid4())
+        url = ctx.file_storage.mk_url(f"artifacts/eval_samples_{artifact_id}.jsonl")
+        self._base = EvaluationArtifactBase(name, url, artifact_id)
+        self.ctx = ctx
+        self.ctx.job.register_artifact(self._base.artifact)
+
+    @property
+    def id(self) -> str:
+        return self._base.id
+
+    @property
+    def name(self) -> str:
+        return self._base.name
+
+    @property
+    def kind(self) -> str:
+        return self._base.kind
+
+    @property
+    def uri(self) -> str:
+        assert self._base.uri is not None
+        return self._base.uri
+
+    def add_samples(self, samples: List[EvalSample]) -> Self:
+        """Add evaluation samples to this artifact.
+
+        Args:
+            samples: List of evaluation samples to add
+
+        Returns:
+            Self for method chaining
+
+        Raises:
+            ValueError: If samples list is empty
+            Exception: If serialization or storage fails
+        """
+        if not samples:
+            raise ValueError("Cannot add empty samples list")
+
+        try:
+            samples_json = self._base.samples_to_adaptive_json(samples)
+            for json_str in samples_json:
+                self.ctx.file_storage.append((json_str + "\n").encode("utf-8"), self.uri)
+            logger.debug(f"Added {len(samples)} samples to artifact {self.id}")
+        except Exception as e:
+            logger.error(f"Failed to add samples to artifact {self.id}: {e}")
+            raise
+
+        return self
+
+    def __repr__(self):
+        return f"EvaluationArtifact(id={self.id}, name={self.name}, kind={self.kind}, uri={self.uri})"

adaptive_harmony/graders/__init__.py

@@ -0,0 +1,20 @@
+from adaptive_harmony import Grade
+
+from .answer_relevancy_judge import AnswerRelevancyGrader
+from .base_grader import BaseGrader
+from .binary_judge import BinaryJudgeGrader
+from .context_relevancy_judge import ContextRelevancyGrader
+from .exceptions import IgnoreScoreException
+from .faithfulness_judge import FaithfulnessGrader
+from .range_judge import RangeJudgeGrader
+
+__all__ = [
+    "BaseGrader",
+    "Grade",
+    "IgnoreScoreException",
+    "BinaryJudgeGrader",
+    "RangeJudgeGrader",
+    "FaithfulnessGrader",
+    "AnswerRelevancyGrader",
+    "ContextRelevancyGrader",
+]

adaptive_harmony/graders/answer_relevancy_judge/__init__.py

@@ -0,0 +1,3 @@
+from .answer_relevancy_judge import AnswerRelevancyGrader
+
+__all__ = ["AnswerRelevancyGrader"]

adaptive_harmony/graders/answer_relevancy_judge/answer_relevancy_judge.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from typing import Literal
+
+import numpy as np
+import pysbd
+from harmony_client import Grade, InferenceModel, StringThread
+from pydantic import BaseModel, Field
+
+from adaptive_harmony.core.utils import stringify_thread
+from adaptive_harmony.graders.answer_relevancy_judge.prompts import DEFAULT_SHOTS, SYSTEM, USER
+from adaptive_harmony.graders.base_grader import BaseGrader
+from adaptive_harmony.graders.faithfulness_judge.faithfulness_judge import SupportedLanguages
+from adaptive_harmony.graders.utils import sample_score_distribution, separate_context_from_last_user_turn
+from adaptive_harmony.logging_table import Table
+
+
+class StatementRelevancy(BaseModel):
+    reason: str = Field(description="The justification for the score given to a statement. Keep it short and concise.")
+    score: Literal[0, 1] = Field(
+        description="The score for the statement. A score of 1 if the statement is relevant to addressing the original input, and 0 if the statement is irrelevant to addressing the input"
+    )
+
+
+class AnswerRelevancyResults(BaseModel):
+    results: list[StatementRelevancy] = Field(description="A list of relevancy results for the statements")
+
+
+class AnswerRelevancyGrader(BaseGrader):
+    def __init__(
+        self,
+        model: InferenceModel,
+        language: SupportedLanguages = "en",
+        grader_key: str = "answer_relevancy_judge",
+        grader_id: str | None = None,
+    ):
+        super().__init__(grader_key)
+        self.model = model
+        self.language = language
+        self.grader_id_or_key = grader_id or grader_key
+        self.sentence_splitter = pysbd.Segmenter(language=language)
+        self.shots = DEFAULT_SHOTS
+
+    async def grade(self, sample: StringThread) -> Grade:
+        _, user_question = separate_context_from_last_user_turn(sample)
+
+        completion = sample.last_content()
+        split_sentences = self.sentence_splitter.segment(completion)
+        sentences = [sentence.strip() for sentence in split_sentences if sentence.strip()]
+        sentences_judge_str = "\n".join(f"{i}: {sentence}" for i, sentence in enumerate(sentences))
+
+        judging_thread = (
+            StringThread()
+            .system(SYSTEM.format(json_schema=self.model.render_schema(AnswerRelevancyResults), shots=self.shots))
+            .user(USER.format(user_question=user_question, statements=sentences_judge_str))
+        )
+
+        try:
+            _, response = await self.model.temperature(0.0).generate_and_validate(
+                judging_thread, AnswerRelevancyResults
+            )
+            results = response.results
+        except Exception as e:
+            self.add_log({"prompt": stringify_thread(judging_thread, sep=f"\n\n{'-' * 10}\n\n"), "error": str(e)})
+            raise
+
+        reason = ""
+        for i, (result, statement) in enumerate(zip(results, sentences)):
+            emoji = "✅" if result.score == 1 else "❌"
+            result = "PASS" if result.score == 1 else "FAIL"
+            statement_display = statement[:150] + ("..." if len(statement) > 150 else "")
+            reason += f"{emoji} Statement {i}: {result}\n Excerpt: {statement_display}:\nReason: {result}\n\n"
+
+        score = np.mean([float(result.score) for result in results]) if results else 0.0
+        self.add_log(
+            {"score": score, "prompt": stringify_thread(judging_thread, sep=f"\n\n{'-' * 10}\n\n"), "reasoning": reason}
+        )
+        return Grade(value=float(score), grader_key=self.grader_id_or_key, reasoning=reason)
+
+    def get_logs(self, clear: bool = False, log_all_samples: bool = False) -> dict[str, float | Table]:
+        # Only clear logs at the end if clear is True
+        logs = super().get_logs(clear=False)
+
+        successfully_scored_samples = [log for log in self._logs if "score" in log]
+
+        # stratified sample range of scores to see high and low
+        if not log_all_samples:
+            subset_successfully_scored_samples = sample_score_distribution(successfully_scored_samples, 15)
+        else:
+            # if we have fewer than 15 samples or we want to log all samples, take them all
+            subset_successfully_scored_samples = successfully_scored_samples
+
+        failed_scored_samples = [log for log in self._logs if "error" in log]
+
+        sample_logs = self.get_sample_tables(subset_successfully_scored_samples, failed_scored_samples)
+
+        logs.update(sample_logs)
+
+        if clear:
+            self.clear_logs()
+
+        return logs