adaptive-harmony 0.1.23__py3-none-any.whl

adaptive_harmony/graders/range_judge/prompts.py

@@ -0,0 +1,232 @@
+import textwrap
+from typing import NamedTuple, TypedDict
+
+from pydantic import BaseModel
+
+from adaptive_harmony import StringThread, StringTurn
+from adaptive_harmony.core.structured_output import render_pydantic_model
+from adaptive_harmony.core.utils import stringify_thread
+from adaptive_harmony.graders.range_judge.types import PromptBuildingBlocks, ReasonedScore
+from adaptive_harmony.graders.utils import (
+    separate_context_from_last_user_turn,
+    validate_thread_last_assistant,
+)
+
+
+class RangeJudgeShot(BaseModel):
+    thread: StringThread
+    reasoning: str
+    score: int
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class RangeShots(TypedDict):
+    reasoning: list[StringTurn]
+    scoring: list[StringTurn]
+
+
+class SubrangeExpectations(NamedTuple):
+    subrange: tuple[int, int]
+    expectation: str
+
+
+COMMON_EVALUATION_PROMPT = """You are an expert evaluator of AI-user interactions.
+You will be given:
+- CONTEXT : previous conversation history, might be empty
+- LAST USER INPUT : the latest input from the user
+- LAST ASSISTANT OUTPUT : the latest output/answer from the AI
+- CRITERIA: the evaluation criteria
+- EVALUATION STEPS : logical reasoning steps to take when evaluating the interaction against the CRITERIA
+"""
+
+
+def get_common_user_template(
+    context: str, last_user_input: str, assistant_answer: str, criteria: str, evaluation_steps: str
+):
+    return f"""CONTEXT\n{context}\n
+LAST USER INPUT\n{last_user_input}\n
+LAST ASSISTANT OUTPUT\n{assistant_answer}\n
+CRITERIA\n{criteria}\n
+EVALUATION STEPS\n{evaluation_steps}"""
+
+
+class RangeScorerTemplates:
+    @staticmethod
+    def get_evaluation_steps(criteria: str) -> StringThread:
+        return (
+            StringThread()
+            .system(
+                textwrap.dedent(
+                    """\
+                    Given an evaluation criteria which outlines how you should judge an interaction between an AI and a user, generate 3-4 concise evaluation steps based on the criteria below.
+
+                    Return your evaluation steps as a numbered list of evaluation steps, such as:
+
+                    Steps list:
+                    1. First step
+                    2. Second step
+                    3. Third step
+                    etc.
+
+                    Focus on specific, concise steps that can be objectively followed based on the evaluation criteria provided.
+                    Don't return any preamble or explanation, only the list.
+                    """
+                )
+            )
+            .user(f"Evaluation Criteria:\n{criteria}\nSteps list:\n")
+        )
+
+    @staticmethod
+    def get_json_reasoned_score_user(
+        context: str, last_user_input: str, assistant_answer: str, criteria: str, evaluation_steps: str
+    ):
+        common = get_common_user_template(context, last_user_input, assistant_answer, criteria, evaluation_steps)
+        return f"{common}\n\nJSON OUTPUT:\n"
+
+    @staticmethod
+    def get_json_reasoned_score(
+        context: str,
+        last_user_input: str,
+        assistant_answer: str,
+        criteria: str,
+        evaluation_steps: str,
+        score_range: tuple[int, int],
+        json_schema: str,
+        subrange_expectations: list[SubrangeExpectations] | None = None,
+        shots: list[StringTurn] | None = None,
+    ) -> StringThread:
+        if subrange_expectations:
+            subrange_expectations_str = "which should correspond to:\n" + "\n".join(
+                [f"{sub.subrange[0]} - {sub.subrange[1]}: {sub.expectation}" for sub in subrange_expectations]
+            )
+        else:
+            subrange_expectations_str = f"where {score_range[1]} indicates strong alignment with the criteria and {score_range[0]} indicates no alignment."
+
+        system_prompt = (
+            COMMON_EVALUATION_PROMPT
+            + f"""
+Your task is to evaluate and score the ASSISTANT ANSWER, strictly following the provided EVALUATION STEPS to evaluate the CRITERIA.
+You must respond with a valid JSON object that matches the following schema:
+
+{json_schema}
+
+Your reasoning for the score:
+- Be specific and grounded in the EVALUATION STEPS
+- Uphold the evaluation objective and nuances expressed in the CRITERIA as the main target
+- Mention specific details, strenghts or shortcomings of the answer, referencing relevant details from the input
+- Be concise, clear, and focused on the evaluation logic.
+- **Never** quote the score itself in the explanation; focus only on reasoning through the evaluation steps
+
+Your final evaluation score must be strictly within the range of [{score_range[0]} - {score_range[1]}], {subrange_expectations_str}
+
+Return only the JSON object after the OUTPUT header, no other text, preamble or explanation.
+"""
+        )
+
+        user_prompt = RangeScorerTemplates.get_json_reasoned_score_user(
+            context, last_user_input, assistant_answer, criteria, evaluation_steps
+        )
+
+        shots = shots or []
+        return StringThread([("system", system_prompt)] + shots + [("user", user_prompt)])
+
+    @staticmethod
+    def get_up_to_score_user(
+        context: str, last_user_input: str, assistant_answer: str, criteria: str, evaluation_steps: str, reasoning: str
+    ) -> str:
+        common = get_common_user_template(context, last_user_input, assistant_answer, criteria, evaluation_steps)
+        return f"{common}\n\nREASONING\n{reasoning}\n\nSCORE: "
+
+    @staticmethod
+    def get_up_to_score(
+        context: str,
+        last_user_input: str,
+        assistant_answer: str,
+        criteria: str,
+        evaluation_steps: str,
+        score_range: tuple[int, int],
+        reasoning: str,
+        subrange_expectations: list[SubrangeExpectations] | None = None,
+        shots: list[StringTurn] | None = None,
+    ) -> StringThread:
+        if subrange_expectations:
+            subrange_expectations_str = "which should correspond to:\n" + "\n".join(
+                [f"{sub.subrange[0]} - {sub.subrange[1]}: {sub.expectation}" for sub in subrange_expectations]
+            )
+        else:
+            subrange_expectations_str = f"where {score_range[1]} indicates strong alignment with the evaluation steps and {score_range[0]} indicates no alignment."
+
+        system_prompt = COMMON_EVALUATION_PROMPT + textwrap.dedent(
+            f"""\
+                - REASONING : the reasoning for the score, following the process described by the EVALUATION STEPS to assess the presented interaction against the CRITERIA
+
+                You must respond only with a score, based on the original CRITERIA and the REASONING for the sample,
+                which should justify your score for the sample.
+
+                Your final evaluation score must be strictly within the range of [{score_range[0]} - {score_range[1]}], {subrange_expectations_str}
+
+                Return only the integer score, nothing before or after.
+                """
+        )
+        user_prompt = RangeScorerTemplates.get_up_to_score_user(
+            context, last_user_input, assistant_answer, criteria, evaluation_steps, reasoning
+        )
+
+        shots = shots or []
+        return StringThread([("system", system_prompt)] + shots + [("user", user_prompt)])
+
+
+def get_prompt_building_blocks(thread: StringThread) -> PromptBuildingBlocks:
+    validate_thread_last_assistant(thread)
+    context_turns, last_user_turn = separate_context_from_last_user_turn(thread)
+    context_str = stringify_thread(StringThread(context_turns))
+    last_assistant_turn = thread.last_content()
+    assert last_user_turn, "There must be at least one user turn"
+    return PromptBuildingBlocks(
+        context=context_str, last_user_turn=last_user_turn, last_assistant_turn=last_assistant_turn
+    )
+
+
+def create_shots(criteria: str, evaluation_steps: str, shots: list[RangeJudgeShot]) -> RangeShots:
+    reasoning_shots: list[StringTurn] = []
+    scoring_shots: list[StringTurn] = []
+    for shot in shots:
+        prompt_components = get_prompt_building_blocks(shot.thread)
+        reasoning_shots.extend(
+            [
+                StringTurn(
+                    role="user",
+                    content=RangeScorerTemplates.get_json_reasoned_score_user(
+                        prompt_components.context,
+                        prompt_components.last_user_turn,
+                        prompt_components.last_assistant_turn,
+                        criteria,
+                        evaluation_steps,
+                    ),
+                ),
+                StringTurn(
+                    role="assistant",
+                    content=render_pydantic_model(ReasonedScore(reasoning=shot.reasoning, score=shot.score)),
+                ),
+            ]
+        )
+        scoring_shots.extend(
+            [
+                StringTurn(
+                    role="user",
+                    content=RangeScorerTemplates.get_up_to_score_user(
+                        prompt_components.context,
+                        prompt_components.last_user_turn,
+                        prompt_components.last_assistant_turn,
+                        criteria,
+                        evaluation_steps,
+                        shot.reasoning,
+                    ),
+                ),
+                StringTurn(role="assistant", content=str(shot.score)),
+            ]
+        )
+
+    return {"reasoning": reasoning_shots, "scoring": scoring_shots}

adaptive_harmony/graders/range_judge/range_judge.py

@@ -0,0 +1,188 @@
+import asyncio
+
+import numpy as np
+
+from adaptive_harmony import InferenceModel, StringThread
+from adaptive_harmony.core.structured_output import JsonParseError
+from adaptive_harmony.core.utils import stringify_thread
+from adaptive_harmony.graders import BaseGrader, Grade
+from adaptive_harmony.graders.range_judge.prompts import (
+    RangeJudgeShot,
+    RangeScorerTemplates,
+    ReasonedScore,
+    SubrangeExpectations,
+    create_shots,
+    get_prompt_building_blocks,
+)
+from adaptive_harmony.graders.utils import (
+    FailedJudgeLog,
+    SuccessJudgeLog,
+    sample_score_distribution,
+)
+from adaptive_harmony.logging_table import Table
+
+
+class RangeJudgeGrader(BaseGrader[SuccessJudgeLog | FailedJudgeLog]):
+    """
+    Scores a thread in a range of integer scores, based on a list of evaluation steps.
+    If evaluation steps are not provided, they are generated from the criteria.
+    The final score is computed as a weighted average of all possible scores,
+    where the weights are the logprobs of each score.
+    You can pass subrange_expectations to the scorer, to help the judge
+    understand the correspondence between score subranges and expected quality levels.
+    """
+
+    def __init__(
+        self,
+        grader_key: str,
+        model: InferenceModel,
+        criteria: str,
+        score_range: tuple[int, int] = (1, 5),
+        evaluation_steps: list[str] | None = None,
+        subrange_expectations: list[SubrangeExpectations] | None = None,
+        shots: list[RangeJudgeShot] | None = None,
+        normalize_score: bool = True,
+    ):
+        model_path: str = model.get_builder_args().get("path")  # type: ignore[assignment]
+        assert model_path.startswith("model_registry://"), "External models cannot be used in RangeJudgeScorer"
+
+        super().__init__(grader_key)
+        self.model = model
+        self.criteria = criteria
+        self.score_range = score_range
+        self.min_score, self.max_score = score_range
+        self.subrange_expectations = subrange_expectations
+        self._shots = shots
+        self.normalize_score = normalize_score
+
+        if evaluation_steps is None:
+            self._str_evaluation_steps = None
+            self._list_eval_steps = None
+        else:
+            self._str_evaluation_steps = "\n".join([f"{i + 1}: {step}" for i, step in enumerate(evaluation_steps)])
+            self._list_eval_steps = evaluation_steps
+
+    @property
+    def evaluation_steps(self) -> list[str] | None:
+        return self._list_eval_steps
+
+    @property
+    def str_evaluation_steps(self) -> str | None:
+        return self._str_evaluation_steps
+
+    @evaluation_steps.setter
+    def evaluation_steps(self, steps: list[str]):
+        self._list_eval_steps = steps
+        self._str_evaluation_steps = "\n".join([f"{i + 1}: {step}" for i, step in enumerate(steps)])
+
+    async def generate_evaluation_steps(self) -> list[str]:
+        thread = await self.model.temperature(0.0).generate(RangeScorerTemplates.get_evaluation_steps(self.criteria))
+        self._str_evaluation_steps = thread.last_content()
+        self._list_eval_steps = self._str_evaluation_steps.split("\n")
+        assert self.evaluation_steps
+        return self.evaluation_steps
+
+    async def grade(self, sample: StringThread) -> Grade:
+        if self.evaluation_steps is None:
+            if self._shots is not None:
+                raise ValueError(
+                    "You cannot pass shots without specifying evaluations steps, since your shots' reasoning must match the steps"
+                )
+            await self.generate_evaluation_steps()
+
+        # Format shots into both user turn formats (for first reasoned scoring step, and for last logprobs scoring step)
+        assert self.str_evaluation_steps
+
+        shots = create_shots(self.criteria, self.str_evaluation_steps, self._shots) if self._shots is not None else {}
+
+        # Separate relevant parts of the prompt turns
+        prompt_components = get_prompt_building_blocks(sample)
+        # Get reasoned scoring thread
+        eval_thread = RangeScorerTemplates.get_json_reasoned_score(
+            context=prompt_components.context,
+            last_user_input=prompt_components.last_user_turn,
+            assistant_answer=prompt_components.last_assistant_turn,
+            criteria=self.criteria,
+            evaluation_steps=self.str_evaluation_steps,
+            score_range=self.score_range,
+            json_schema=self.model.render_schema(ReasonedScore),
+            subrange_expectations=self.subrange_expectations,
+            shots=shots.get("reasoning"),
+        )
+        eval_str_prompt = stringify_thread(eval_thread, sep=f"\n\n{'-' * 10}\n\n")
+        try:
+            _, reasoned_score = await self.model.temperature(0.0).generate_and_validate(eval_thread, ReasonedScore)
+        except JsonParseError as e:
+            self.add_log({"prompt": eval_str_prompt, "error": f"{str(e)}\n\nCOMPLETION:\n{e.completion}"})
+            raise
+        except Exception as e:
+            self.add_log({"prompt": eval_str_prompt, "error": str(e)})
+            raise
+
+        # Get a prompt that includes the reasoning for the sample, all the way to form-filling the score
+        up_to_score_thread = RangeScorerTemplates.get_up_to_score(
+            context=prompt_components.context,
+            last_user_input=prompt_components.last_user_turn,
+            assistant_answer=prompt_components.last_assistant_turn,
+            criteria=self.criteria,
+            evaluation_steps=self.str_evaluation_steps,
+            score_range=self.score_range,
+            reasoning=reasoned_score.reasoning,
+            subrange_expectations=self.subrange_expectations,
+            shots=shots.get("scoring"),
+        )
+
+        # Get logprobs for each possible final score
+        possible_score_ints = [s for s in range(self.min_score, self.max_score + 1)]
+        logprobs = await asyncio.gather(
+            *[self.model.temperature(0.0).logprobs(up_to_score_thread.assistant(f"{s}")) for s in possible_score_ints]
+        )
+
+        # Convert to probabilities and compute weighted average
+        probs = np.exp(logprobs - np.logaddexp.reduce(logprobs))
+        weighted_score = np.average(possible_score_ints, weights=probs)
+
+        final_score: float = weighted_score
+        if self.normalize_score:  # normalize to 0-1 range
+            final_score = (weighted_score - self.min_score) / (self.max_score - self.min_score)
+
+        str_prompt = stringify_thread(eval_thread, sep=f"\n\n{'-' * 10}\n\n")
+        self.add_log({"score": final_score, "prompt": str_prompt, "reasoning": reasoned_score.reasoning})
+
+        metadata = dict(
+            criteria=self.criteria,
+            raw_avg_score=float(weighted_score),
+            scale_range=(self.min_score, self.max_score),
+            score_probabilities={str(score): float(prob) for score, prob in zip(possible_score_ints, probs)},
+            evaluation_steps=self.evaluation_steps,
+            reasoning=reasoned_score.reasoning,
+        )
+
+        return Grade(value=float(final_score), grader_key=self.grader_key, reasoning=reasoned_score.reasoning)
+
+    def add_log(self, log_data: SuccessJudgeLog | FailedJudgeLog) -> None:
+        self._logs.append(log_data)
+
+    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

adaptive_harmony/graders/range_judge/types.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel, Field
+
+
+class ReasonedScore(BaseModel):
+    reasoning: str = Field(description="String reasoning to support the rationale behind the score")
+    score: int = Field(description="Integer score for the interaction, must be within the specified score range")
+
+
+class PromptBuildingBlocks(BaseModel):
+    context: str
+    last_user_turn: str
+    last_assistant_turn: str

adaptive_harmony/graders/reward_server_grader.py

@@ -0,0 +1,36 @@
+import asyncio
+
+from adaptive_harmony import StringThread
+from adaptive_harmony.core.reward_client.client import Request, RewardClient, Turn
+from adaptive_harmony.graders.base_grader import BaseGrader, Grade
+
+
+class RewardServerGrader(BaseGrader):
+    def __init__(self, grader_key: str, grader_id: str, reward_server_ip: str):
+        super().__init__(grader_key)
+        self.reward_client = RewardClient(reward_server_ip)
+        self.grader_id_or_key = grader_id or grader_key
+        self._setup_task = None
+        self._setup_lock = None
+
+    async def _ensure_setup(self):
+        if self._setup_lock is None:
+            self._setup_lock = asyncio.Lock()
+
+        if self._setup_task is None:
+            async with self._setup_lock:
+                if self._setup_task is None:
+                    self._setup_task = asyncio.create_task(self.reward_client.setup())
+
+        await self._setup_task
+
+    async def grade(self, sample: StringThread) -> Grade:
+        await self._ensure_setup()
+
+        response = await self.reward_client.score(
+            Request(
+                turns=[Turn(content=turn.content, role=turn.role) for turn in sample.get_turns()],
+                metadata=sample.metadata,
+            )
+        )
+        return Grade(value=response.reward, grader_key=self.grader_id_or_key, reasoning=response.metadata.get("reason"))