Upload sharegpt_polar.py with huggingface_hub

sharegpt_polar.py

@@ -0,0 +1,462 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, Iterable
+
+from datasets import Dataset, load_dataset
+import httpx
+
+import verifiers as vf
+from verifiers.types import Messages, State
+
+DEFAULT_MODEL = "internlm/POLAR-7B"
+POOL_ENDPOINT = "/pooling"
+
+
+def _ensure_messages(conversations: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
+    messages: list[dict[str, Any]] = []
+    for turn in conversations:
+        role = turn.get("from") or turn.get("role")
+        content = turn.get("value") or turn.get("content")
+        if role == "system":
+            messages.append({"role": "system", "content": content})
+        elif role == "human" or role == "user":
+            messages.append({"role": "user", "content": content})
+        elif role in {"gpt", "assistant"}:
+            messages.append({"role": "assistant", "content": content})
+    return messages
+
+
+def _has_assistant(conversations: Iterable[dict[str, Any]]) -> bool:
+    return any(
+        (turn.get("from") or turn.get("role")) in {"gpt", "assistant"}
+        for turn in conversations
+    )
+
+
+def _partition_conversation(
+    messages: list[dict[str, Any]]
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[list[dict[str, Any]]], list[str]]:
+    assistant_indices = [idx for idx, msg in enumerate(messages) if msg["role"] == "assistant"]
+    if not assistant_indices:
+        raise ValueError("Conversation must include at least one assistant response")
+
+    first_assistant_idx = assistant_indices[0]
+    prompt_messages = messages[:first_assistant_idx]
+    if not any(msg["role"] == "user" for msg in prompt_messages):
+        raise ValueError("Conversation must include a user message before the first assistant turn")
+
+    reference_messages = [messages[idx] for idx in assistant_indices]
+
+    future_turns: list[list[dict[str, Any]]] = []
+    user_contexts: list[str] = []
+    assistant_indices_with_end = assistant_indices + [len(messages)]
+    for current_idx, next_idx in zip(assistant_indices, assistant_indices_with_end[1:]):
+        env_msgs: list[dict[str, Any]] = []
+        user_context_lines: list[str] = []
+        for i in range(current_idx + 1, next_idx):
+            turn = messages[i]
+            role = turn["role"]
+            content = turn["content"]
+            if role == "system":
+                continue
+            if role == "user":
+                line = (content or "").strip()
+                if line:
+                    user_context_lines.append(line)
+                    env_msgs.append({"role": "user", "content": line})
+            else:
+                env_msgs.append(turn)
+        future_turns.append(env_msgs)
+        user_contexts.append("\n".join(user_context_lines).strip())
+
+    return prompt_messages, reference_messages, future_turns, user_contexts
+
+
+def _extract_last_value(value: Any) -> float | None:
+    current: Any = value
+    while isinstance(current, list) and current:
+        current = current[-1]
+    if isinstance(current, (int, float)):
+        return float(current)
+    return None
+
+
+class PoolingClient:
+    def __init__(
+        self,
+        base_url: str,
+        model: str = DEFAULT_MODEL,
+        timeout: float = 30.0,
+        logger: logging.Logger | None = None,
+        enable_logging: bool = False,
+    ):
+        self.base_url = base_url.rstrip("/")
+        if not self.base_url.startswith("http"):
+            self.base_url = f"https://{self.base_url}"
+        self.timeout = timeout
+        self.model = model
+        self.logger = logger or logging.getLogger("sharegpt_polar.PoolingClient")
+        self.enable_logging = enable_logging
+
+    @staticmethod
+    def encode(sample: dict[str, Any]) -> str:
+        def _messages_to_text(messages: list[dict[str, Any]] | None) -> str:
+            if not messages:
+                return ""
+            return "\n".join(msg.get("content", "") for msg in messages if msg.get("content"))
+
+        prompt_text = _messages_to_text(sample.get("prompt"))
+        reference_text = _messages_to_text(sample.get("reference"))
+        output_text = _messages_to_text(sample.get("output"))
+
+        reference_cat = f"{prompt_text}\n{reference_text}" if reference_text else prompt_text
+        output_cat = f"{prompt_text}\n{output_text}" if output_text else prompt_text
+
+        return f"{reference_cat}<|reward|>{output_cat}[UNUSED_TOKEN_130]"
+
+    async def score(self, payload: list[dict[str, Any]]) -> dict[str, Any] | list[Any]:
+        encoded_payload = [self.encode(item) for item in payload]
+        if self.enable_logging:
+            self.logger.debug(
+                "Sending reward request",
+                extra={
+                    "payload_size": len(payload),
+                    "model": self.model,
+                    "endpoint": self.base_url,
+                },
+            )
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.post(
+                f"{self.base_url}{POOL_ENDPOINT}",
+                json={"model": self.model, "input": encoded_payload},
+                headers={"Content-Type": "application/json"},
+            )
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as exc:
+            if self.enable_logging:
+                self.logger.error(
+                    "Reward server request failed",
+                    extra={
+                        "status": exc.response.status_code,
+                        "body": exc.response.text,
+                    },
+                )
+            raise RuntimeError(
+                f"Pooling request failed: {exc.response.status_code} {exc.response.text}"
+            ) from exc
+        if self.enable_logging:
+            self.logger.debug(
+                "Received reward response",
+                extra={
+                    "status": response.status_code,
+                    "model": self.model,
+                },
+            )
+        return response.json()
+
+
+async def polar_reward(
+    prompt: Messages,
+    completion: Messages,
+    info: dict[str, Any],
+    reward_client: PoolingClient,
+    logger: logging.Logger | None = None,
+    enable_logging: bool = False,
+    **_: Any,
+) -> float:
+    assistant_turns = [msg for msg in completion if msg.get("role") == "assistant"]
+    if not assistant_turns:
+        if enable_logging:
+            (logger or logging.getLogger("sharegpt_polar.reward")).debug(
+                "No assistant turn available for reward",
+                extra={"prompt": prompt, "completion": completion},
+            )
+        return 0.0
+
+    payload = [
+        {
+            "prompt": prompt,
+            "reference": info.get("reference", []),
+            "output": [assistant_turns[-1]],
+        }
+    ]
+    try:
+        data = await reward_client.score(payload)
+    except RuntimeError as err:
+        if enable_logging:
+            (logger or logging.getLogger("sharegpt_polar.reward")).exception(
+                "Reward request failed", extra={"error": str(err), "payload": payload}
+            )
+        raise
+    if enable_logging:
+        (logger or logging.getLogger("sharegpt_polar.reward")).debug(
+            "Reward response received", extra={"response": data}
+        )
+    if isinstance(data, dict):
+        if "data" in data:
+            scores = data["data"][0]["data"]
+            last_value = _extract_last_value(scores)
+            if last_value is not None:
+                return last_value
+        if "rewards" in data and data["rewards"]:
+            last_value = _extract_last_value(data["rewards"])
+            if last_value is not None:
+                return last_value
+    if isinstance(data, list) and data:
+        last_value = _extract_last_value(data)
+        if last_value is not None:
+            return last_value
+    if enable_logging:
+        (logger or logging.getLogger("sharegpt_polar.reward")).error(
+            "Unexpected reward payload", extra={"response": data}
+        )
+    raise RuntimeError(f"Unexpected reward response: {data}")
+
+
+class ShareGPTPolarEnv(vf.MultiTurnEnv):
+    def __init__(
+        self,
+        dataset: Dataset,
+        rubric: vf.Rubric,
+        *,
+        enable_logging: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(dataset=dataset, rubric=rubric, **kwargs)
+        self.enable_logging = enable_logging
+        self.logger = logging.getLogger("sharegpt_polar.env")
+
+    async def setup_state(self, state: State, **kwargs: Any) -> State:
+        state.setdefault("future_turns", state["info"].get("future_turns", []))
+        return state
+
+    async def is_completed(self, messages: Messages, state: State, **kwargs: Any) -> bool:
+        total_turns = len(state["info"].get("reference", []))
+        if self.enable_logging:
+            self.logger.debug(
+                "Checking completion state",
+                extra={"current_turn": state.get("turn", 0), "total_turns": total_turns},
+            )
+        return state.get("turn", 0) >= total_turns
+
+    async def env_response(self, messages: Messages, state: State, **kwargs: Any) -> tuple[Messages, State]:
+        future_turns: list[list[dict[str, Any]]] = state.get("future_turns", [])
+        turn_index = state.get("turn", 0) - 1
+        if self.enable_logging:
+            self.logger.debug(
+                "Providing future turn",
+                extra={"turn_index": turn_index, "future_turn_count": len(future_turns)},
+            )
+        if 0 <= turn_index < len(future_turns):
+            return future_turns[turn_index], state
+        return [], state
+
+    def process_chat_format_vllm(  # type: ignore[override]
+        self,
+        prompt: list[dict[str, Any]],
+        completion: list[dict[str, Any]],
+        state: State,
+        processing_class: Any,
+        mask_env_responses: bool = False,
+    ) -> tuple[list[int], list[int], list[int], list[int], list[float]]:
+        # Clean messages to remove tool-related fields that might trigger template errors
+        def clean_message(msg: dict[str, Any]) -> dict[str, Any]:
+            return {k: v for k, v in msg.items() if k not in {"tool_calls", "tool_call_id"}}
+        
+        responses = state.get("responses", [])
+        responses_idx = 0
+        zipped: list[tuple[dict[str, Any], Any | None]] = []
+        for turn in completion:
+            if turn.get("role") == "assistant":
+                zipped.append((turn, responses[responses_idx]))
+                responses_idx += 1
+            else:
+                zipped.append((turn, None))
+        assert len(responses) == responses_idx, "Responses not fully consumed"
+        assert len(zipped) == len(completion), "Length mismatch"
+
+        clean_prompt = [clean_message(msg) for msg in prompt]
+        prompt_ids: list[int] = processing_class.apply_chat_template(
+            conversation=clean_prompt,  # type: ignore[arg-type]
+            add_generation_prompt=True,
+            tools=None,
+        )
+        messages_consumed = [clean_message(m) for m in prompt]
+        prompt_mask: list[int] = [0] * len(prompt_ids)
+        completion_ids: list[int] = []
+        completion_mask: list[int] = []
+        completion_logprobs: list[float] = []
+        i = 0
+        while i < len(zipped):
+            message, response = zipped[i]
+            clean_msg = clean_message(message)
+            if message.get("role") == "assistant":
+                if response is not None:
+                    completion_turn_ids = self.parse_chat_completion_tokens(response)
+                    completion_turn_mask = [1] * len(completion_turn_ids)
+                    completion_turn_logprobs = self.parse_chat_completion_logprobs(response)
+                else:
+                    completion_turn_ids = []
+                    completion_turn_mask = []
+                    completion_turn_logprobs = []
+                completion_ids.extend(completion_turn_ids)
+                completion_mask.extend(completion_turn_mask)
+                completion_logprobs.extend(completion_turn_logprobs)
+                messages_consumed.append(clean_msg)
+                i += 1
+                continue
+
+            consecutive_messages = [clean_msg]
+            j = i + 1
+            while j < len(zipped) and zipped[j][0].get("role") != "assistant":
+                consecutive_messages.append(clean_message(zipped[j][0]))
+                j += 1
+
+            base_tokens: list[int] = processing_class.apply_chat_template(
+                conversation=messages_consumed,  # type: ignore[arg-type]
+                add_generation_prompt=True,
+                tools=None,
+            )
+            extended_tokens: list[int] = processing_class.apply_chat_template(
+                conversation=messages_consumed + consecutive_messages,  # type: ignore[arg-type]
+                add_generation_prompt=True,
+                tools=None,
+            )
+            prefix_len = 0
+            max_len = min(len(base_tokens), len(extended_tokens))
+            while prefix_len < max_len and base_tokens[prefix_len] == extended_tokens[prefix_len]:
+                prefix_len += 1
+            if self.enable_logging and prefix_len != len(base_tokens):
+                self.logger.debug(
+                    "Token prefix adjusted",
+                    extra={"prefix_len": prefix_len, "base_len": len(base_tokens)},
+                )
+            completion_turn_ids = extended_tokens[prefix_len:]
+            if mask_env_responses:
+                completion_turn_mask = [0] * len(completion_turn_ids)
+            else:
+                completion_turn_mask = [1] * len(completion_turn_ids)
+            completion_turn_logprobs = [0.0] * len(completion_turn_ids)
+            completion_ids.extend(completion_turn_ids)
+            completion_mask.extend(completion_turn_mask)
+            completion_logprobs.extend(completion_turn_logprobs)
+            messages_consumed.extend(consecutive_messages)
+            i = j
+
+        return (
+            prompt_ids,
+            prompt_mask,
+            completion_ids,
+            completion_mask,
+            completion_logprobs,
+        )
+
+
+def load_environment(
+    dataset_name: str | None = None,
+    *,
+    dataset_split: str = "train",
+    dataset_files: dict[str, str] | None = None,
+    data_path: str | Path | None = None,
+    server_address: str,
+    reward_model: str = DEFAULT_MODEL,
+    reward_scheme: type[vf.Rubric] | None = None,
+    max_turns: int = -1,
+    enable_logging: bool = False,
+    logger: logging.Logger | None = None,
+    **env_kwargs: Any,
+) -> ShareGPTPolarEnv:
+    if dataset_name is None and data_path is None:
+        raise ValueError("Either 'dataset_name' or 'data_path' must be provided")
+
+    if dataset_name is not None:
+        hf_dataset = load_dataset(dataset_name, split=dataset_split, data_files=dataset_files)
+    else:
+        hf_dataset = load_dataset("json", data_files=str(data_path), split="train")
+
+    def to_multi_turn(example: dict[str, Any]) -> dict[str, Any]:
+        conversations = example.get("conversations") or []
+        if not _has_assistant(conversations):
+            return {
+                "prompt": [],
+                "info": {
+                    "reference": [],
+                    "future_turns": [],
+                },
+                "valid": False,
+            }
+
+        messages = _ensure_messages(conversations)
+        prompt, reference, future_turns, user_contexts = _partition_conversation(messages)
+        flattened_future = [msg for block in future_turns for msg in block]
+        if any(msg.get("role") != "user" for msg in flattened_future):
+            return {
+                "prompt": [],
+                "info": {
+                    "reference": [],
+                    "future_turns": [],
+                },
+                "valid": False,
+            }
+        if any(not msg.get("content") for msg in flattened_future):
+            return {
+                "prompt": [],
+                "info": {
+                    "reference": [],
+                    "future_turns": [],
+                },
+                "valid": False,
+            }
+
+        return {
+            "prompt": prompt,
+            "info": {
+                "reference": reference,
+                "future_turns": future_turns,
+                "user_contexts": user_contexts,
+            },
+            "valid": True,
+        }
+
+    dataset = hf_dataset.map(to_multi_turn, remove_columns=hf_dataset.column_names)
+    dataset = dataset.filter(lambda example: example.get("valid", False))
+    if "valid" in dataset.column_names:
+        dataset = dataset.remove_columns("valid")
+
+    effective_logger = logger or logging.getLogger("sharegpt_polar")
+    if enable_logging:
+        effective_logger.info(
+            "Initializing ShareGPTPolar environment",
+            extra={
+                "dataset_name": dataset_name,
+                "data_path": str(data_path) if data_path else None,
+                "server_address": server_address,
+            },
+        )
+
+    client = PoolingClient(
+        base_url=server_address,
+        model=reward_model,
+        logger=effective_logger,
+        enable_logging=enable_logging,
+    )
+
+    rubric_cls = reward_scheme or vf.Rubric
+    rubric = rubric_cls(funcs=[polar_reward])
+    rubric.class_objects["reward_client"] = client
+    rubric.class_objects["logger"] = effective_logger
+    rubric.class_objects["enable_logging"] = enable_logging
+
+    env_kwargs.setdefault("max_concurrent", 1)
+    return ShareGPTPolarEnv(
+        dataset=dataset,
+        rubric=rubric,
+        max_turns=max_turns,
+        enable_logging=enable_logging,
+        **env_kwargs,
+    )
+
+