Upload folder using huggingface_hub

pyproject.toml

@@ -0,0 +1,23 @@
+[project]
+name = "roleplay-selfplay"
+version = "0.1.0"
+description = "Roleplay self-play environment with CAI critique-and-revision"
+tags = ["multi-turn", "roleplay", "self-play", "train", "eval"]
+requires-python = ">=3.10"
+dependencies = [
+    "verifiers>=0.1.6",
+    "datasets",
+    "httpx",
+    "huggingface_hub",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+include = ["roleplay_selfplay.py", "pyproject.toml"]
+
+[tool.verifiers.eval]
+num_examples = 5
+rollouts_per_example = 1

roleplay_selfplay.py

@@ -0,0 +1,466 @@
+"""
+Roleplay Self-Play RL Environment (MultiTurnEnv-based).
+
+Architecture:
+1. Dataset provides system prompt + first user turn (the "seed")
+2. Standard rollout loop: model generates Character A (assistant), env generates Character B (user)
+3. All turns judged by external judge model for reward signal
+4. Model only trained on Character A turns (standard GRPO with logprobs)
+"""
+
+import asyncio
+import logging
+import random
+import re
+from typing import Any
+
+import httpx
+from datasets import load_dataset
+from openai import AsyncOpenAI
+
+import verifiers as vf
+from verifiers.envs.multiturn_env import MultiTurnEnv
+from verifiers.types import Info, Messages, SamplingArgs, State
+
+logger = logging.getLogger("roleplay_selfplay")
+
+
+# ---------------------------------------------------------------------------
+# Utility helpers
+# ---------------------------------------------------------------------------
+
+
+def strip_think_tags(text: str) -> str:
+    if not text:
+        return text
+    cleaned = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL | re.IGNORECASE)
+    cleaned = re.sub(r"<think>.*$", "", cleaned, flags=re.DOTALL | re.IGNORECASE)
+    return cleaned.strip()
+
+
+def count_words(text: str) -> int:
+    if not text:
+        return 0
+    cleaned = re.sub(r"```.*?```", "", text, flags=re.DOTALL)
+    return len([w for w in re.split(r"\s+", cleaned.strip()) if w])
+
+
+def detect_lists(text: str) -> tuple[bool, dict[str, Any]]:
+    if not text:
+        return False, {}
+    text_clean = re.sub(r"```.*?```", "", text, flags=re.DOTALL)
+    bullet_items = re.findall(r"^\s*[-*\u2022\u25cf]\s+.+$", text_clean, re.MULTILINE)
+    numbered_items = re.findall(r"^\s*\d+[.)]\s+.+$", text_clean, re.MULTILINE)
+    has = len(bullet_items) >= 3 or len(numbered_items) >= 3
+    return has, {"bullet_count": len(bullet_items), "numbered_count": len(numbered_items)}
+
+
+def detect_structured_markers(text: str) -> tuple[bool, dict[str, Any]]:
+    if not text:
+        return False, {}
+    markers = 0
+    examples: list[str] = []
+    xml = re.findall(r"<([a-zA-Z_]\w{2,})>.*?</\1>", text, re.DOTALL)
+    if xml:
+        markers += 1
+        examples.append(f"XML: {xml[0]}")
+    brackets = re.findall(r"\[([A-Z_]{3,})\]", text)
+    if brackets:
+        markers += 1
+        examples.append(f"Bracket: [{brackets[0]}]")
+    return markers > 0, {"total_markers": markers, "examples": examples}
+
+
+# ---------------------------------------------------------------------------
+# Turn-level judge prompt (symmetric -- no user/assistant labels)
+# ---------------------------------------------------------------------------
+
+TURN_JUDGE_PROMPT = """You are evaluating a single turn from a roleplay conversation.
+Score ONLY the turn inside <evaluate_turn> tags. The context is for reference.
+
+<conversation_context>
+{context}
+</conversation_context>
+
+<evaluate_turn>
+{turn_text}
+</evaluate_turn>
+
+How good is this turn as a piece of collaborative roleplay writing?
+
+IMPORTANT -- Length is NOT quality:
+- A short, punchy turn that lands perfectly is BETTER than a long, meandering one.
+- Do NOT give higher scores to longer turns.
+
+Score 1-9:
+1-2 = Great -- pulls you in, makes the scene better
+3-4 = Solid -- does its job well
+5-6 = Mediocre -- functional but flat or generic
+7-8 = Weak -- filler, repetitive, or disconnected
+9 = Broken -- incoherent, off-topic, or format garbage
+
+<rationale>[1-2 sentences]</rationale>
+<score>[1-9]</score>"""
+
+
+# ---------------------------------------------------------------------------
+# Generation prompt for Character B (used in env_response)
+# ---------------------------------------------------------------------------
+
+GENERATE_USER_TURN_SYSTEM = """You are playing a CHARACTER in a roleplay conversation.
+You are NOT an AI assistant. You are the other character in this scene.
+
+Rules:
+- You are a CHARACTER, not a helpful assistant. Act like a person in this scene.
+- Have your own goals, reactions, emotions, and agency
+- You can disagree, push back, be surprised, get angry, be curious, etc.
+- Advance the scene -- introduce complications, reveal information, change dynamics
+- Use vivid sensory and emotional detail
+- Write natural narrative prose (no bullet points, no structured markers)
+- Keep the turn roughly 50-200 words
+- Write ONLY your character's turn, nothing else"""
+
+
+# ---------------------------------------------------------------------------
+# Environment
+# ---------------------------------------------------------------------------
+
+
+class RoleplaySelfPlayEnv(MultiTurnEnv):
+    """
+    Roleplay self-play environment using MultiTurnEnv for training compatibility.
+
+    The standard rollout loop generates Character A (assistant) turns with logprobs.
+    env_response() generates Character B (user) turns via separate API calls.
+    An external judge model scores all turns for the reward signal.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        judge_client: AsyncOpenAI,
+        judge_model: str,
+        judge_temperature: float = 0.8,
+        judge_min_p: float = 0.05,
+        selfplay_min_turns: int = 4,
+        selfplay_max_turns: int = 10,
+        list_penalty_threshold: float = 0.5,
+        list_penalty_multiplier: float = 0.1,
+        **kwargs,
+    ):
+        # Placeholder rubric; replaced below with bound method
+        super().__init__(max_turns=-1, dataset=dataset, rubric=vf.Rubric(), **kwargs)
+
+        self.judge_client = judge_client
+        self.judge_model = judge_model
+        self.judge_temperature = judge_temperature
+        self.judge_min_p = judge_min_p
+        self.selfplay_min_turns = selfplay_min_turns
+        self.selfplay_max_turns = selfplay_max_turns
+        self.list_penalty_threshold = list_penalty_threshold
+        self.list_penalty_multiplier = list_penalty_multiplier
+
+        # Replace rubric with bound method that accesses self.judge_client directly
+        self.rubric = vf.Rubric(funcs=[self._judge_reward], weights=[1.0])
+
+        # Set by rollout() for use in env_response()
+        self._rollout_client: AsyncOpenAI | None = None
+        self._rollout_model: str = ""
+
+    async def _judge_reward(
+        self, completion: Messages, state: State, prompt: Messages, **kwargs,
+    ) -> float:
+        """Judge all conversation turns and return mean reward."""
+        # Reconstruct flat turn list from prompt + completion
+        turns: list[dict[str, str]] = []
+        for msg in (prompt if isinstance(prompt, list) else []):
+            if isinstance(msg, dict) and msg.get("role") not in ("system", None):
+                turns.append({"content": msg["content"], "side": "B"})
+        for msg in (completion if isinstance(completion, list) else []):
+            if isinstance(msg, dict):
+                side = "A" if msg.get("role") == "assistant" else "B"
+                turns.append({"content": msg["content"], "side": side})
+
+        if not turns:
+            return 0.0
+
+        # Judge all turns concurrently
+        tasks = [
+            _judge_single_turn(
+                self.judge_client, self.judge_model,
+                self.judge_temperature, self.judge_min_p,
+                i, t["content"], turns,
+            )
+            for i, t in enumerate(turns)
+        ]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        rewards = [r if not isinstance(r, Exception) else 0.0 for r in results]
+        mean_reward = sum(rewards) / len(rewards) if rewards else 0.0
+
+        # Global list penalty
+        turns_with_lists = sum(1 for t in turns if detect_lists(t["content"])[0])
+        list_ratio = turns_with_lists / len(turns) if turns else 0
+        if list_ratio > self.list_penalty_threshold:
+            mean_reward *= self.list_penalty_multiplier
+
+        return mean_reward
+
+    async def setup_state(self, state: State, **kwargs) -> State:
+        # Random total conversation turns, then compute assistant turn count
+        total = random.randint(self.selfplay_min_turns, self.selfplay_max_turns)
+        # state["turn"] counts assistant (model) turns only
+        # total includes seed user turn: 1 seed + N assistant + (N-1) user = 2N
+        # So assistant turns = total // 2
+        state["target_model_turns"] = max(2, total // 2)
+        return state
+
+    async def is_completed(self, messages: Messages, state: State, **kwargs) -> bool:
+        if await self.prompt_too_long(state):
+            return True
+        return state["turn"] >= state.get("target_model_turns", 3)
+
+    async def env_response(
+        self, messages: Messages, state: State, **kwargs
+    ) -> tuple[Messages, State]:
+        """Generate Character B (user) turn using the same model."""
+        client = self._rollout_client
+        model = self._rollout_model
+
+        # Extract system prompt from the original prompt
+        system_prompt = ""
+        for msg in state["prompt"]:
+            if isinstance(msg, dict) and msg.get("role") == "system":
+                system_prompt = msg["content"]
+                break
+
+        # Build conversation context for Character B generation
+        convo_text = f"[Roleplay System Prompt]: {system_prompt}\n\n"
+        turn_idx = 0
+        for msg in state["prompt"]:
+            if isinstance(msg, dict) and msg.get("role") != "system":
+                turn_idx += 1
+                convo_text += f"[Turn {turn_idx}]:\n{msg['content']}\n\n"
+        for msg in state["completion"]:
+            if isinstance(msg, dict):
+                turn_idx += 1
+                convo_text += f"[Turn {turn_idx}]:\n{msg['content']}\n\n"
+
+        gen_messages: list[dict[str, str]] = [
+            {"role": "system", "content": GENERATE_USER_TURN_SYSTEM},
+            {
+                "role": "user",
+                "content": (
+                    f"Here is the conversation so far:\n\n{convo_text}\n\n"
+                    "Write the next turn:"
+                ),
+            },
+        ]
+
+        try:
+            response = await client.chat.completions.create(
+                model=model,
+                messages=gen_messages,  # type: ignore
+                temperature=0.9,
+                max_tokens=1024,
+            )
+            turn_text = strip_think_tags(response.choices[0].message.content or "")
+        except Exception as e:
+            logger.error(f"Failed to generate Character B turn: {e}")
+            turn_text = "*continues the scene*"
+
+        if not turn_text or len(turn_text.strip()) < 5:
+            turn_text = "*continues the scene*"
+
+        return [{"role": "user", "content": turn_text}], state
+
+    async def rollout(
+        self,
+        client: AsyncOpenAI,
+        model: str,
+        prompt: Messages,
+        completion: Messages | None = None,
+        answer: str = "",
+        state: State = {},
+        task: str = "default",
+        info: Info | None = None,
+        example_id: int = 0,
+        sampling_args: SamplingArgs | None = None,
+        **kwargs,
+    ) -> tuple[Messages, State]:
+        """Capture client/model for env_response, then delegate to MultiTurnEnv loop."""
+        self._rollout_client = client
+        self._rollout_model = model
+        return await super().rollout(
+            client, model, prompt, completion, answer,
+            state, task, info, example_id, sampling_args, **kwargs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Judging helper
+# ---------------------------------------------------------------------------
+
+
+async def _judge_single_turn(
+    judge_client: AsyncOpenAI,
+    judge_model: str,
+    judge_temperature: float,
+    judge_min_p: float,
+    turn_index: int,
+    turn_text: str,
+    all_turns: list[dict[str, str]],
+) -> float:
+    """Judge a single turn and return a reward in [0, 1]."""
+    context = ""
+    for j in range(turn_index):
+        context += f"[Turn {j + 1}]:\n{all_turns[j]['content']}\n\n"
+
+    prompt_text = TURN_JUDGE_PROMPT.format(
+        context=context.strip() if context.strip() else "(This is the opening turn)",
+        turn_text=turn_text,
+    )
+
+    try:
+        response = await judge_client.chat.completions.create(
+            model=judge_model,
+            messages=[{"role": "user", "content": prompt_text}],  # type: ignore
+            temperature=judge_temperature,
+            max_tokens=2048,
+            extra_body={"min_p": judge_min_p} if judge_min_p else {},
+        )
+        judge_text = response.choices[0].message.content or ""
+
+        score_match = re.search(
+            r"<score>\s*([1-9])\s*</score>", judge_text, re.IGNORECASE,
+        )
+        if score_match:
+            score = int(score_match.group(1))
+        else:
+            fallback = re.search(r"<score>\s*([1-9])", judge_text, re.IGNORECASE)
+            score = int(fallback.group(1)) if fallback else 9
+
+        reward = 1.0 - (score / 10.0)
+
+        # Length bias control
+        wc = count_words(turn_text)
+        if wc < 15:
+            reward *= 0.3
+        elif wc < 40:
+            reward *= 0.7 + 0.3 * ((wc - 15) / 25)
+        elif wc <= 250:
+            pass
+        elif wc <= 400:
+            reward *= 1.0 - 0.4 * ((wc - 250) / 150)
+        else:
+            reward *= max(0.2, 0.6 - 0.4 * ((wc - 400) / 200))
+
+        # Structured marker penalty
+        has_markers, _ = detect_structured_markers(turn_text)
+        if has_markers:
+            reward = 0.0
+
+        return reward
+    except Exception as e:
+        logger.error(f"Judge failed for turn {turn_index}: {e}")
+        return 0.0
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+
+def load_environment(
+    dataset_name: str = "Delta-Vector/Hydrus-UnsafeRLHF",
+    dataset_split: str = "train",
+    judge_model: str = "default",
+    judge_base_url: str = "https://reduces-bone-benefit-endangered.trycloudflare.com/v1",
+    judge_temperature: float = 0.8,
+    judge_min_p: float = 0.05,
+    judge_timeout: float = 1200.0,
+    max_concurrent_scoring: int = 32,
+    selfplay_min_turns: int = 4,
+    selfplay_max_turns: int = 10,
+    list_penalty_threshold: float = 0.5,
+    list_penalty_multiplier: float = 0.1,
+    **kwargs,
+) -> RoleplaySelfPlayEnv:
+    """Load the roleplay self-play environment."""
+
+    dataset = load_dataset(dataset_name, split=dataset_split)
+
+    def transform_example(example: dict, idx: int) -> dict:
+        system_prompt = ""
+        first_user_msg = ""
+
+        if "conversations" in example:
+            for c in example["conversations"]:
+                role = c.get("from", c.get("role", ""))
+                content = c.get("value", c.get("content", ""))
+                if role == "system":
+                    system_prompt = content
+                elif role in ("human", "user") and not first_user_msg:
+                    first_user_msg = content
+                    break
+        elif "prompt" in example:
+            prompt_content = example["prompt"]
+            if isinstance(prompt_content, str):
+                first_user_msg = prompt_content
+            elif isinstance(prompt_content, list):
+                for m in prompt_content:
+                    if m.get("role") == "system":
+                        system_prompt = m.get("content", "")
+                    elif m.get("role") == "user" and not first_user_msg:
+                        first_user_msg = m.get("content", "")
+        else:
+            raise ValueError("Dataset must have 'prompt' or 'conversations' column")
+
+        if not system_prompt:
+            system_prompt = "You are a creative and engaging roleplay partner."
+        if not first_user_msg:
+            first_user_msg = "*A stranger approaches.*"
+
+        messages: list[dict[str, str]] = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": first_user_msg},
+        ]
+
+        return {
+            "prompt": messages,
+            "info": {
+                "original_system_prompt": system_prompt,
+                "seed_user_turn": first_user_msg,
+            },
+        }
+
+    columns_to_remove = list(dataset.column_names)
+    dataset = dataset.map(
+        transform_example, with_indices=True, remove_columns=columns_to_remove,
+    )
+
+    http_client = httpx.AsyncClient(
+        limits=httpx.Limits(
+            max_connections=max_concurrent_scoring,
+            max_keepalive_connections=max_concurrent_scoring,
+        ),
+        timeout=judge_timeout,
+    )
+    judge_client = AsyncOpenAI(
+        base_url=judge_base_url,
+        api_key="dummy-key",
+        http_client=http_client,
+    )
+
+    return RoleplaySelfPlayEnv(
+        dataset=dataset,
+        judge_client=judge_client,
+        judge_model=judge_model,
+        judge_temperature=judge_temperature,
+        judge_min_p=judge_min_p,
+        selfplay_min_turns=selfplay_min_turns,
+        selfplay_max_turns=selfplay_max_turns,
+        list_penalty_threshold=list_penalty_threshold,
+        list_penalty_multiplier=list_penalty_multiplier,
+        **kwargs,
+    )