adaptive-harmony 0.1.23__tar.gz

adaptive_harmony-0.1.23/PKG-INFO

@@ -0,0 +1,37 @@
+Metadata-Version: 2.4
+Name: adaptive-harmony
+Version: 0.1.23
+Summary: Adaptive Harmony training recipes and utilities for LLM fine-tuning
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Python: >=3.12
+Requires-Dist: harmony-client
+Requires-Dist: rich>=13.7.0
+Requires-Dist: datasets>=2.14.0
+Requires-Dist: hf-xet>=1.1.2
+Requires-Dist: loguru>=0.7.2
+Requires-Dist: pydantic-settings>=2.9.1
+Requires-Dist: pydantic>=2.10.5
+Requires-Dist: pybars3>=0.9.7
+Requires-Dist: pysbd>=0.3.4
+Requires-Dist: websockets>=15.0.1
+Requires-Dist: setproctitle>=1.3.3
+Requires-Dist: pydantic-xml>=2.16.0
+Requires-Dist: openai>=1.42.0
+Requires-Dist: pillow>=11.3.0
+Requires-Dist: boto3>=1.40
+Requires-Dist: tomli>=2
+Provides-Extra: tensorboard
+Requires-Dist: tensorboardX>=2.6.2.2; extra == "tensorboard"
+Provides-Extra: mlflow
+Requires-Dist: mlflow>=3.4.0; extra == "mlflow"
+Provides-Extra: wandb
+Requires-Dist: wandb>=0.19.11; extra == "wandb"
+Provides-Extra: all-monitoring
+Requires-Dist: tensorboardX>=2.6.2.2; extra == "all-monitoring"
+Requires-Dist: mlflow>=3.4.0; extra == "all-monitoring"
+Requires-Dist: wandb>=0.19.11; extra == "all-monitoring"
+Provides-Extra: mlflow-skinny
+Requires-Dist: mlflow-skinny>=3.4.0; extra == "mlflow-skinny"

adaptive_harmony-0.1.23/pyproject.toml

@@ -0,0 +1,110 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+version = "0.1.23"
+name = "adaptive-harmony"
+description = "Adaptive Harmony training recipes and utilities for LLM fine-tuning"
+requires-python = ">=3.12"
+classifiers = [
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: Implementation :: CPython",
+    "Programming Language :: Python :: Implementation :: PyPy",
+]
+dependencies = [
+    "harmony-client",
+    "rich>=13.7.0",
+    "datasets>=2.14.0",
+    "hf-xet>=1.1.2",
+    "loguru>=0.7.2",
+    "pydantic-settings>=2.9.1",
+    "pydantic>=2.10.5",
+    "pybars3>=0.9.7",
+    "pysbd>=0.3.4",
+    "websockets>=15.0.1",
+    "setproctitle>=1.3.3",
+    "pydantic-xml>=2.16.0",
+    "openai>=1.42.0",
+    "pillow>=11.3.0",
+    "boto3>=1.40",
+    "tomli>=2",
+]
+
+[project.optional-dependencies]
+tensorboard = ["tensorboardX>=2.6.2.2"]
+mlflow = ["mlflow>=3.4.0"]
+wandb = ["wandb>=0.19.11"]
+all_monitoring = ["tensorboardX>=2.6.2.2", "mlflow>=3.4.0", "wandb>=0.19.11"]
+mlflow_skinny = ["mlflow-skinny>=3.4.0"]
+
+[dependency-groups]
+dev = ["pyright>=1.1.407", "pytest", "pytest-asyncio>=0.23.5", "ruff>=0.14.3"]
+
+[tool.pyright]
+pythonVersion = "3.12"
+pythonPlatform = "Linux"
+typeCheckingMode = "standard"
+reportUntypedFunctionDecorator = "warning"
+reportMissingTypeStubs = "warning"
+reportMissingImports = false
+reportUnusedVariable = "warning"
+reportUnusedImport = "warning"
+reportDuplicateImport = "error"
+reportArgumentType = "error"
+reportPossiblyUnboundVariable = "warning"
+reportIncompatibleMethodOverride = "warning"
+exclude = [
+    ".venv/**",
+    "**/__pycache__",
+    "tests",
+    "examples",
+    "python/adaptive_harmony/metric_logger.py",
+    "python/adaptive_harmony/runtime",          # auto-generated code
+]
+extraPaths = ["python"]
+
+[tool.ruff]
+line-length = 120
+lint.ignore = [
+    "E731",
+    "E402",
+    "UP035",
+    "UP045",
+    "UP040",
+    "UP006",
+    "F841",
+    "E501",
+]
+lint.select = [
+    "E",     # pycodestyle errors
+    "F",     # pyflakes
+    "I",     # isort (import sorting)
+    "UP",    # pyupgrade (modern Python syntax)
+    "ASYNC", # async-specific checks
+]
+lint.fixable = ["I", "UP", "F401"]
+exclude = [
+    ".venv/**",
+    "**/__pycache__",
+    "tests",
+    "examples",
+    "python/adaptive_harmony/runtime", # auto-generated code
+]
+
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+markers = [
+    "harmony: mark tests that require harmony to run (deselect with '-m \"not harmony\"'')",
+]
+
+[tool.setuptools]
+package-dir = { "" = "python" }
+
+[tool.setuptools.packages.find]
+where = ["python"]
+
+[tool.uv]
+cache-keys = [{ file = "pyproject.toml" }]

adaptive_harmony-0.1.23/python/adaptive_harmony/__init__.py

@@ -0,0 +1,162 @@
+# ruff: noqa: F403, F401
+from typing import TYPE_CHECKING
+
+from harmony_client import (
+    EvalSample as EvalSample,
+)
+from harmony_client import (
+    EvalSampleInteraction as EvalSampleInteraction,
+)
+from harmony_client import (
+    Grade as Grade,
+)
+from harmony_client import (
+    HarmonyClient as HarmonyClient,
+)
+from harmony_client import (
+    HarmonyJobNotifier as HarmonyJobNotifier,
+)
+from harmony_client import (
+    InferenceModel as InferenceModel,
+)
+from harmony_client import (
+    JobArtifact as JobArtifact,
+)
+from harmony_client import (
+    JobNotifier as JobNotifier,
+)
+from harmony_client import (
+    ModelBuilder as ModelBuilder,
+)
+from harmony_client import (
+    StageNotifier as StageNotifier,
+)
+from harmony_client import (
+    StringThread as StringThread,
+)
+from harmony_client import (
+    TokenizedThread as TokenizedThread,
+)
+from harmony_client import (
+    TrainingModel as TrainingModel,
+)
+from harmony_client import (
+    get_client as get_client,
+)
+from harmony_client import parameters as parameters
+from harmony_client import runtime as runtime
+from rich.progress import Progress
+
+if TYPE_CHECKING:
+    from harmony_client import StringTurn as StringTurn
+else:
+    from typing import NamedTuple
+
+    class StringTurn(NamedTuple):
+        role: str
+        content: str
+
+
+from harmony_client.artifacts.custom_artifact import CustomArtifact
+from harmony_client.artifacts.dataset_artifact import DatasetArtifact
+from harmony_client.file_storage import (
+    FileStorage,
+    FileStorageConfig,
+    LocalFileStorageConfig,
+    S3FileStorageConfig,
+    StoredFile,
+)
+
+import adaptive_harmony.core.rl_utils as rl_utils
+from adaptive_harmony.core.dataset import DataSet
+from adaptive_harmony.core.schedulers import CombinedSchedule, CosineScheduler, CosineSchedulerWithoutWarmup, Scheduler
+from adaptive_harmony.evaluation.evaluation_artifact import EvaluationArtifact
+from adaptive_harmony.metric_logger import Logger, WandbLogger
+
+# Ensure key classes are available at module level
+__all__ = [
+    "StringThread",
+    "StringTurn",
+    "TokenizedThread",
+    "InferenceModel",
+    "ModelBuilder",
+    "TrainingModel",
+    "HarmonyClient",
+    "get_client",
+    "DataSet",
+    "CosineScheduler",
+    "CombinedSchedule",
+    "CosineSchedulerWithoutWarmup",
+    "Scheduler",
+    "WandbLogger",
+    "Logger",
+    "FileStorage",
+    "FileStorageConfig",
+    "LocalFileStorageConfig",
+    "S3FileStorageConfig",
+    "StoredFile",
+    "EvaluationArtifact",
+    "CustomArtifact",
+    "DatasetArtifact",
+    "rl_utils",
+    "Grade",
+    "EvalSample",
+    "EvalSampleInteraction",
+    "JobArtifact",
+]
+
+
+# Patch StringThread to use rich for display
+from harmony_client.runtime.model_artifact_save import save_with_artifact
+
+from adaptive_harmony.core.display import _stringthread_repr, _tokenizedthread_repr
+from adaptive_harmony.core.image_utils import string_thread_to_html_string
+
+# Patch InferenceModel to have json output capabilities
+from adaptive_harmony.core.structured_output import generate_and_validate, render_pydantic_model, render_schema
+
+StringThread.__repr__ = _stringthread_repr  # type: ignore
+TokenizedThread.__repr__ = _tokenizedthread_repr  # type: ignore
+setattr(StringThread, "_repr_html_", string_thread_to_html_string)
+setattr(InferenceModel, "generate_and_validate", generate_and_validate)
+setattr(InferenceModel, "render_schema", staticmethod(render_schema))
+setattr(InferenceModel, "render_pydantic_model", staticmethod(render_pydantic_model))
+
+_original_training_model_save = TrainingModel.save
+
+
+async def _save_with_artifact_wrapper(model: TrainingModel, model_name: str, inference_only: bool = True, ctx=None):
+    return await save_with_artifact(model, model_name, inference_only, ctx, _original_training_model_save)
+
+
+setattr(TrainingModel, "save", _save_with_artifact_wrapper)
+
+
+async def spawn_train(self: ModelBuilder, name: str, max_batch_size: int) -> TrainingModel:
+    fut = await self.spawn_train_with_progress(name, max_batch_size)  # type:ignore
+
+    with Progress() as pbar:
+        task = pbar.add_task("Loading model", total=1000)
+
+        while (prog := await fut._await_progress()) != 1.0:
+            pbar.update(task, completed=prog, total=1.0)
+        pbar.update(task, completed=1.0, total=1.0)
+
+    return await fut.get()
+
+
+async def spawn_inference(self: ModelBuilder, name: str) -> InferenceModel:
+    fut = await self.spawn_inference_with_progress(name)  # type:ignore
+
+    with Progress() as pbar:
+        task = pbar.add_task("Loading model", total=1000)
+
+        while (prog := await fut._await_progress()) != 1.0:
+            pbar.update(task, completed=prog, total=1.0)
+        pbar.update(task, completed=1.0, total=1.0)
+
+    return await fut.get()
+
+
+setattr(ModelBuilder, "spawn_inference", spawn_inference)
+setattr(ModelBuilder, "spawn_train", spawn_train)

adaptive_harmony-0.1.23/python/adaptive_harmony/common/__init__.py

@@ -0,0 +1,40 @@
+from .callbacks import (
+    CheckpointCallback as CheckpointCallback,
+)
+from .callbacks import (
+    EnvironmentValidationCallback as EnvironmentValidationCallback,
+)
+from .callbacks import (
+    GenerateSamplesCallback as GenerateSamplesCallback,
+)
+from .callbacks import (
+    GraderEvalCallback as GraderEvalCallback,
+)
+from .callbacks import (
+    RecipeCallback as RecipeCallback,
+)
+from .callbacks import (
+    ValidationLossCallback as ValidationLossCallback,
+)
+from .dpo import DPO as DPO
+from .env_grpo import ENVGRPO
+from .grpo import GRPO as GRPO
+from .gspo import GSPO as GSPO
+from .ppo import PPO as PPO
+from .rm import RewardModelling as RewardModelling
+from .sft import SFT as SFT
+
+__all__ = [
+    "SFT",
+    "PPO",
+    "GRPO",
+    "ENVGRPO",
+    "DPO",
+    "RewardModelling",
+    "RecipeCallback",
+    "GenerateSamplesCallback",
+    "ValidationLossCallback",
+    "CheckpointCallback",
+    "GraderEvalCallback",
+    "EnvironmentValidationCallback",
+]

adaptive_harmony-0.1.23/python/adaptive_harmony/common/callbacks.py

@@ -0,0 +1,219 @@
+from abc import abstractmethod
+from typing import Any
+
+import numpy as np
+from harmony_client import (
+    InferenceModel,
+    StringThread,
+    TrainingModel,
+)
+from loguru import logger
+
+from adaptive_harmony.core.utils import async_map, async_map_fallible
+from adaptive_harmony.environment import EnvironmentFactory
+from adaptive_harmony.graders import BaseGrader
+from adaptive_harmony.logging_table import Table
+
+
+class RecipeCallback:
+    def __init__(self, frequency: float, log_key_prefix: str | None = None):
+        self.frequency = frequency
+        self.last_call = -1.0
+        self.log_key_prefix = log_key_prefix
+
+    async def maybe_call(self, current_percentage: float) -> dict[str, Any]:
+        if current_percentage - self.last_call >= self.frequency:
+            self.last_call = current_percentage
+            callback_dict = await self.callback(current_percentage)
+            prefixed_dict = {
+                (f"{self.log_key_prefix}/{key}" if self.log_key_prefix else key): value
+                for key, value in callback_dict.items()
+            }
+            return prefixed_dict
+        return {}
+
+    @abstractmethod
+    async def callback(self, current_percentage: float) -> dict[str, Any]: ...
+
+
+class GenerateSamplesCallback(RecipeCallback):
+    def __init__(
+        self,
+        thread_set: list[StringThread],
+        model: InferenceModel,
+        frequency: float,
+        log_key: str = "samples",
+    ):
+        super().__init__(frequency, log_key_prefix="generation")
+        self.thread_set = thread_set
+        self.model = model
+        self.log_key = log_key
+
+    async def callback(self, current_percentage: float) -> dict[str, Any]:
+        logger.info("Entering generation callback...")
+        generation_tokens = await async_map_fallible(self.model.generate_tokens, self.thread_set)
+        generation_results = await async_map_fallible(self.model.detokenize_thread, generation_tokens)
+        gen_lengths = [sample.len_last_turn() for sample in generation_tokens]
+
+        generation_logs = {
+            self.log_key: Table()
+            .add_column(
+                "system",
+                [
+                    sample.get_turns()[0].content if sample.get_turns()[0].role == "system" else ""
+                    for sample in generation_results
+                ],
+            )
+            .add_column(
+                "prompt",
+                [
+                    repr(
+                        StringThread(
+                            sample.get_turns()[1:-1]
+                            if (sample.get_turns() and sample.get_turns()[0].role == "system")
+                            else sample.get_turns()[:-1]
+                        )
+                    )
+                    for sample in generation_results
+                ],
+            )
+            .add_column("response", [response.last_content() for response in generation_results]),
+            "generation_length_mean": np.mean(gen_lengths).item(),
+            "generation_length_std": np.std(gen_lengths).item(),
+            "num_samples": len(generation_results),
+        }
+        return generation_logs
+
+
+class ValidationLossCallback(RecipeCallback):
+    def __init__(
+        self,
+        validation_set: list[StringThread],
+        model: InferenceModel,
+        frequency: float = 0.1,
+        log_key: str = "loss",
+    ):
+        super().__init__(frequency, log_key_prefix="validation")
+        self.validation_set = validation_set
+        self.model = model
+        self.log_key = log_key
+
+    async def callback(self, current_percentage: float) -> dict[str, float]:
+        logger.info("Entering validation loss callback...")
+        losses = []
+        tokens = await async_map_fallible(self.model.tokenize_thread, self.validation_set)
+        logprobs = await async_map(self.model.logprobs_per_token, tokens)
+        losses = [-(sum(lp) / len(lp)) for lp in logprobs]
+
+        return {self.log_key: sum(losses) / len(losses)}
+
+
+class CheckpointCallback(RecipeCallback):
+    def __init__(
+        self,
+        model: TrainingModel,
+        checkpoint_name: str,
+        frequency: float = 0.2,
+    ):
+        super().__init__(frequency, log_key_prefix="checkpointing")
+        self.last_call = 0.0  # avoid saving the model at the first period
+        self.model = model
+        self.model_log_name = checkpoint_name
+
+    async def callback(self, current_percentage: float):
+        logger.info(f"Saving checkpoint at {current_percentage * 100} % of training ...")
+        await self.model.save(f"{self.model_log_name}-{round(current_percentage, 3)}")
+        return {}
+
+
+class GraderEvalCallback(RecipeCallback):
+    def __init__(
+        self,
+        validation_set: list[StringThread],
+        model: InferenceModel,
+        grader: BaseGrader,
+        frequency: float,
+        log_key: str = "validation",
+        clear_grader_logs: bool = True,
+        temperature: float = 0.0,
+    ):
+        super().__init__(frequency, log_key_prefix=log_key)
+        self.validation_set = validation_set
+        self.model = model
+        self.grader = grader
+        self.clear_grader_logs = clear_grader_logs
+        self.temperature = temperature
+
+    async def callback(self, current_percentage: float) -> dict[str, float | Table]:
+        logger.info("Entering grader evaluation callback...")
+        temp_model = self.model.temperature(self.temperature)
+
+        tokenized_results = await async_map_fallible(temp_model.generate_tokens, self.validation_set)
+        string_results = await async_map(temp_model.detokenize_thread, tokenized_results)
+        grades = await async_map_fallible(self.grader.grade, string_results)
+        gen_lengths = [sample.len_last_turn() for sample in tokenized_results]
+
+        grader_logs = self.grader.get_logs(clear=self.clear_grader_logs)
+        return {
+            **{f"rewards/{key}": value for key, value in grader_logs.items()},
+            "generation_length_mean": float(np.mean(gen_lengths).item()),
+            "generation_length_std": float(np.std(gen_lengths).item()),
+            "num_samples": float(len(grades)),
+        }
+
+
+class EnvironmentValidationCallback(RecipeCallback):
+    def __init__(
+        self,
+        validation_set: list[StringThread],
+        model: InferenceModel,
+        env_factory: EnvironmentFactory,
+        frequency: float,
+        log_key: str = "validation",
+        clear_env_logs: bool = True,
+        temperature: float = 0.0,
+        num_samples_log: int = 0,
+    ):
+        super().__init__(frequency, log_key_prefix=log_key)
+        self.validation_set = validation_set
+        self.model = model
+        self.env_factory = env_factory
+        self.clear_env_logs = clear_env_logs
+        self.temperature = temperature
+        self.num_samples_log = num_samples_log
+
+    async def generate_trajectory(self, initial_thread: StringThread) -> tuple[StringThread, float, int]:
+        env = self.env_factory.create_environment(initial_thread.metadata)
+        temp_model = self.model.temperature(self.temperature)
+        trajectory, trajectory_score = await env.generate_trajectory_and_grade(temp_model, initial_thread)
+        num_turns = len([turn for turn in trajectory.get_turns() if turn.role == "assistant"])
+        return trajectory, trajectory_score.cumulative_score, num_turns
+
+    async def callback(self, current_percentage: float) -> dict[str, float | Table]:
+        logger.info("Entering environment validation callback...")
+
+        results = await async_map_fallible(self.generate_trajectory, self.validation_set)
+
+        trajectories = [traj for traj, _, _ in results]
+        scores = [score for _, score, _ in results]
+        num_turns_list = [num_turns for _, _, num_turns in results]
+
+        validation_logs = {
+            "score_mean": np.mean(scores).item(),
+            "score_std": np.std(scores).item(),
+            "num_turns_mean": np.mean(num_turns_list).item(),
+            "num_turns_std": np.std(num_turns_list).item(),
+            "num_samples": len(results),
+        }
+
+        env_logs = self.env_factory.get_logs(clear=self.clear_env_logs)
+        validation_logs.update({f"env/{key}": value for key, value in env_logs.items()})
+
+        if self.num_samples_log > 0:
+            samples = [repr(traj) for traj in trajectories[: self.num_samples_log]]
+            samples_scores = scores[: self.num_samples_log]
+            table = Table().add_column("trajectory", samples).add_column("score", samples_scores)
+            validation_logs["samples"] = table
+
+        logger.info(f"Validation Mean score: {validation_logs['score_mean']:.4f}")
+        return validation_logs