mirror of https://github.com/razor-ai/soup.git
100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
"""Turn a stream of :class:`Trace` events into preference pairs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Iterable, Iterator
|
|
|
|
from soup_cli.data.traces.parsers import Trace
|
|
|
|
SUPPORTED_SIGNALS: tuple[str, ...] = (
|
|
"thumbs_up",
|
|
"regenerations",
|
|
"user_edit",
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PreferencePair:
|
|
prompt: str
|
|
chosen: str
|
|
rejected: str
|
|
source: str = ""
|
|
|
|
def to_jsonl_dict(self) -> dict:
|
|
return {
|
|
"prompt": self.prompt,
|
|
"chosen": self.chosen,
|
|
"rejected": self.rejected,
|
|
"source": self.source,
|
|
}
|
|
|
|
|
|
def _group_by_prompt(traces: Iterable[Trace]) -> dict[str, list[Trace]]:
|
|
grouped: dict[str, list[Trace]] = {}
|
|
for trace in traces:
|
|
grouped.setdefault(trace.prompt, []).append(trace)
|
|
return grouped
|
|
|
|
|
|
def _pair_thumbs(traces: Iterable[Trace]) -> Iterator[PreferencePair]:
|
|
grouped = _group_by_prompt(traces)
|
|
for prompt, entries in grouped.items():
|
|
ups = [t for t in entries if t.signal == "thumbs_up"]
|
|
downs = [t for t in entries if t.signal == "thumbs_down"]
|
|
if not ups or not downs:
|
|
continue
|
|
yield PreferencePair(
|
|
prompt=prompt, chosen=ups[0].output,
|
|
rejected=downs[0].output, source="thumbs_up",
|
|
)
|
|
|
|
|
|
def _pair_regenerations(traces: Iterable[Trace]) -> Iterator[PreferencePair]:
|
|
grouped = _group_by_prompt(traces)
|
|
for prompt, entries in grouped.items():
|
|
relevant = [
|
|
t for t in entries
|
|
if t.signal in ("regenerated", "thumbs_up", "thumbs_down")
|
|
or t.regen_order is not None
|
|
]
|
|
if len(relevant) < 2:
|
|
continue
|
|
ordered = sorted(relevant, key=lambda t: t.regen_order or 0)
|
|
if ordered[0].output == ordered[-1].output:
|
|
continue
|
|
yield PreferencePair(
|
|
prompt=prompt, chosen=ordered[-1].output,
|
|
rejected=ordered[0].output, source="regenerations",
|
|
)
|
|
|
|
|
|
def _pair_user_edit(traces: Iterable[Trace]) -> Iterator[PreferencePair]:
|
|
for trace in traces:
|
|
if trace.signal != "user_edit" or not trace.edited_output:
|
|
continue
|
|
if trace.output == trace.edited_output:
|
|
continue
|
|
yield PreferencePair(
|
|
prompt=trace.prompt, chosen=trace.edited_output,
|
|
rejected=trace.output, source="user_edit",
|
|
)
|
|
|
|
|
|
def build_pairs(
|
|
traces: Iterable[Trace], *, signal: str,
|
|
) -> Iterator[PreferencePair]:
|
|
"""Produce preference pairs from ``traces`` according to ``signal``."""
|
|
if signal not in SUPPORTED_SIGNALS:
|
|
raise ValueError(
|
|
f"unknown signal '{signal}'. "
|
|
f"Supported: {', '.join(SUPPORTED_SIGNALS)}"
|
|
)
|
|
traces = list(traces) # materialise for multi-pass signals
|
|
if signal == "thumbs_up":
|
|
yield from _pair_thumbs(traces)
|
|
elif signal == "regenerations":
|
|
yield from _pair_regenerations(traces)
|
|
elif signal == "user_edit":
|
|
yield from _pair_user_edit(traces)
|