[tts] Refactor inference and test-time scaling ownership

Author: taivu1998

experiments/evals/run_reasoning_tts.py

@@ -10,12 +10,12 @@
 
 from rigging.log_setup import configure_logging
 
-from marin.evaluation.evaluators.evaluator import ModelConfig
+from marin.inference.chat_completions import OpenAIChatCompletionProvider
+from marin.inference.model_config import ModelConfig
 from marin.inference.vllm_server import VllmEnvironment
 from marin.test_time_scaling import (
     DEFAULT_REASONING_SELECTORS,
     CandidateGenerationConfig,
-    OpenAIChatCompletionProvider,
     SelectorName,
     TestTimeScalingConfig,
     build_run_summary,

lib/marin/src/marin/evaluation/evaluators/evalchemy_evaluator.py

@@ -36,8 +36,9 @@
 from rigging.filesystem import filesystem as marin_filesystem
 
 from marin.evaluation.evaluation_config import WANDB_PROJECT, EvalTaskConfig
-from marin.evaluation.evaluators.evaluator import Evaluator, ModelConfig, launch_evaluate_with_ray
+from marin.evaluation.evaluators.evaluator import Evaluator, launch_evaluate_with_ray
 from marin.evaluation.utils import is_remote_path, upload_to_gcs
+from marin.inference.model_config import ModelConfig
 from marin.inference.vllm_server import resolve_model_name_or_path
 
 logger = logging.getLogger(__name__)

lib/marin/src/marin/evaluation/evaluators/evaluator.py

@@ -4,11 +4,11 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any
 
 from fray.v1.cluster import Entrypoint, EnvironmentConfig, JobRequest, ResourceConfig, current_cluster
 
 from marin.evaluation.evaluation_config import EvalTaskConfig
+from marin.inference.model_config import ModelConfig
 from marin.utils import remove_tpu_lockfile_on_exit
 from rigging.log_setup import configure_logging as _init_logging
 
@@ -27,35 +27,6 @@ def __str__(self):
         return f"{self.name}=={self.version}" if self.version else self.name
 
 
-@dataclass
-class ModelConfig:
-    name: str
-    """The name of the model e.g., allenai/olmo-7b"""
-
-    path: str | None
-    """
-    The path to the model checkpoint. Can be a local path or a path on GCS.
-    """
-
-    engine_kwargs: dict[str, Any]
-    """
-    Additional keyword arguments to pass to the vLLM engine.
-    """
-
-    generation_params: dict | None = None
-    """
-    Additional keyword arguments passed to the SamplingParams for the vLLM engine
-    """
-
-    apply_chat_template: bool = False
-    """
-    Whether or not this model was trained with a Chat Template in the tokenizer
-    """
-
-    base_eval_run_name: str | None = None
-    """Custom base name for wandb runs."""
-
-
 class Evaluator(ABC):
     @abstractmethod
     def launch_evaluate_with_ray(

lib/marin/src/marin/evaluation/evaluators/harbor_evaluator.py

@@ -28,8 +28,9 @@
 from rigging.filesystem import open_url
 
 from marin.evaluation.evaluation_config import EvalTaskConfig
-from marin.evaluation.evaluators.evaluator import Evaluator, ModelConfig, launch_evaluate_with_ray
+from marin.evaluation.evaluators.evaluator import Evaluator, launch_evaluate_with_ray
 from marin.evaluation.utils import download_from_gcs, is_remote_path, upload_to_gcs
+from marin.inference.model_config import ModelConfig
 from marin.inference.vllm_server import VLLM_NATIVE_PIP_PACKAGES, VllmEnvironment, resolve_vllm_mode
 from marin.utils import fsspec_exists, fsspec_glob
 

lib/marin/src/marin/evaluation/evaluators/levanter_lm_eval_evaluator.py

@@ -15,10 +15,10 @@
 from levanter.tracker.wandb import WandbConfig
 from levanter.trainer import TrainerConfig
 
+from fray.v1.cluster.ray.deps import build_runtime_env_for_packages
 from marin.evaluation.evaluation_config import EvalTaskConfig, convert_to_levanter_task_config
-from marin.evaluation.evaluators.evaluator import ModelConfig
 from marin.evaluation.evaluators.levanter_tpu_evaluator import LevanterTpuEvaluator
-from fray.v1.cluster.ray.deps import build_runtime_env_for_packages
+from marin.inference.model_config import ModelConfig
 
 logger = logging.getLogger(__name__)
 

lib/marin/src/marin/evaluation/evaluators/levanter_tpu_evaluator.py

@@ -6,7 +6,8 @@
 from fray.v1.cluster import ResourceConfig
 
 from marin.evaluation.evaluation_config import EvalTaskConfig
-from marin.evaluation.evaluators.evaluator import Evaluator, ModelConfig, launch_evaluate_with_ray
+from marin.evaluation.evaluators.evaluator import Evaluator, launch_evaluate_with_ray
+from marin.inference.model_config import ModelConfig
 
 
 class LevanterTpuEvaluator(Evaluator, ABC):

lib/marin/src/marin/evaluation/evaluators/lm_evaluation_harness_evaluator.py

@@ -15,8 +15,9 @@
 from rigging.filesystem import open_url, url_to_fs
 
 from marin.evaluation.evaluation_config import EvalTaskConfig
-from marin.evaluation.evaluators.evaluator import Evaluator, ModelConfig, launch_evaluate_with_ray
+from marin.evaluation.evaluators.evaluator import Evaluator, launch_evaluate_with_ray
 from marin.evaluation.utils import is_remote_path, upload_to_gcs
+from marin.inference.model_config import ModelConfig
 from marin.inference.vllm_server import VLLM_NATIVE_PIP_PACKAGES, VllmEnvironment, resolve_vllm_mode
 
 logger = logging.getLogger(__name__)

lib/marin/src/marin/evaluation/evaluators/simple_evaluator.py

@@ -9,7 +9,8 @@
 from fray.v1.cluster import ResourceConfig
 
 from marin.evaluation.evaluation_config import EvalTaskConfig
-from marin.evaluation.evaluators.evaluator import Evaluator, ModelConfig, launch_evaluate_with_ray
+from marin.evaluation.evaluators.evaluator import Evaluator, launch_evaluate_with_ray
+from marin.inference.model_config import ModelConfig
 from marin.inference.vllm_server import VLLM_NATIVE_PIP_PACKAGES, resolve_model_name_or_path, resolve_vllm_mode
 
 

lib/marin/src/marin/evaluation/run.py

@@ -21,13 +21,14 @@
 from fray.v1.cluster import TpuConfig as V1TpuConfig
 
 from marin.evaluation.evaluation_config import EvaluationConfig
-from marin.evaluation.evaluators.evaluator import Evaluator, ModelConfig
+from marin.evaluation.evaluators.evaluator import Evaluator
 from marin.evaluation.evaluators.evalchemy_evaluator import EvalchemyEvaluator
 from marin.evaluation.evaluators.harbor_evaluator import HarborEvaluator
 from marin.evaluation.evaluators.levanter_lm_eval_evaluator import LevanterLmEvalEvaluator
 from marin.evaluation.evaluators.lm_evaluation_harness_evaluator import LMEvaluationHarnessEvaluator
 from marin.evaluation.evaluators.simple_evaluator import SimpleEvaluator
 from marin.evaluation.utils import discover_hf_checkpoints
+from marin.inference.model_config import ModelConfig
 from marin.utils import fsspec_exists
 
 logger = logging.getLogger(__name__)

lib/marin/src/marin/inference/chat_completions.py

@@ -0,0 +1,74 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from openai import OpenAI
+from openai.types.chat import ChatCompletion
+
+
+@dataclass(frozen=True)
+class ChatCompletionRequest:
+    """OpenAI-compatible chat completion request parameters."""
+
+    messages: tuple[dict[str, str], ...]
+    num_completions: int
+    temperature: float
+    top_p: float = 1.0
+    max_tokens: int | None = None
+    seed: int | None = None
+    logprobs: bool = False
+
+    def __post_init__(self) -> None:
+        if self.num_completions <= 0:
+            raise ValueError("num_completions must be positive")
+        if self.temperature < 0:
+            raise ValueError("temperature must be non-negative")
+        if not 0 < self.top_p <= 1.0:
+            raise ValueError("top_p must be in the interval (0, 1]")
+        if self.max_tokens is not None and self.max_tokens <= 0:
+            raise ValueError("max_tokens must be positive when set")
+
+
+class CompletionProvider(Protocol):
+    """Protocol for chat completion backends used by inference clients."""
+
+    def complete_messages(self, request: ChatCompletionRequest) -> ChatCompletion:
+        """Return an OpenAI-compatible chat completion response."""
+
+
+class OpenAIChatCompletionProvider:
+    """Minimal synchronous OpenAI-compatible completion provider."""
+
+    def __init__(
+        self,
+        *,
+        server_url: str,
+        model: str,
+        api_key: str = "marin-tts",
+        timeout: float | None = None,
+        extra_request_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        self._client = OpenAI(base_url=server_url, api_key=api_key, timeout=timeout)
+        self._model = model
+        self._extra_request_kwargs = dict(extra_request_kwargs or {})
+
+    def complete_messages(self, request: ChatCompletionRequest) -> ChatCompletion:
+        request_kwargs: dict[str, Any] = {
+            "model": self._model,
+            "messages": list(request.messages),
+            "n": request.num_completions,
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "logprobs": request.logprobs,
+            **self._extra_request_kwargs,
+        }
+        if request.max_tokens is not None:
+            request_kwargs["max_tokens"] = request.max_tokens
+        if request.seed is not None:
+            request_kwargs["seed"] = request.seed
+
+        return self._client.chat.completions.create(**request_kwargs)