Upload folder using huggingface_hub

build/lib/sharegpt_polar.py

@@ -0,0 +1,113 @@
+"""ShareGPT + POLAR reward environment."""
+
+from __future__ import annotations
+
+import json
+from http import HTTPStatus
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset, load_dataset
+import httpx
+
+import verifiers as vf
+from verifiers.types import Messages
+
+
+DEFAULT_SERVER = "wealth-intent-submissions-range.trycloudflare.com"
+DEFAULT_MODEL = "internlm/POLAR-7B"
+
+
+def _load_sharegpt_dataset(path: str | Path) -> Dataset:
+    dataset = load_dataset("json", data_files=str(path), split="train")
+
+    def to_single_turn(example: dict[str, Any]) -> dict[str, Any]:
+        human_turn = next(
+            turn["value"] for turn in example["conversations"] if turn["from"] == "human"
+        )
+        assistant_turn = next(
+            turn["value"] for turn in example["conversations"] if turn["from"] == "gpt"
+        )
+        return {
+            "prompt": [{"role": "user", "content": human_turn}],
+            "info": {
+                "reference": [{"role": "assistant", "content": assistant_turn}],
+            },
+        }
+
+    return dataset.map(to_single_turn, remove_columns=dataset.column_names)
+
+
+async def polar_reward(
+    prompt: Messages,
+    completion: Messages,
+    info: dict[str, Any],
+    reward_client: "PolarClient",
+    **_: Any,
+) -> float:
+    assistant_turns = [msg for msg in completion if msg.get("role") == "assistant"]
+    if not assistant_turns:
+        return 0.0
+
+    payload = [
+        {
+            "prompt": prompt,
+            "reference": info.get("reference", []),
+            "output": [assistant_turns[-1]],
+        }
+    ]
+    scores = await reward_client.score(payload)
+    return float(scores[0]) if scores else 0.0
+
+
+def load_environment(
+    data_path: str | Path,
+    *,
+    server_address: str = DEFAULT_SERVER,
+    reward_model: str = DEFAULT_MODEL,
+    reward_scheme: type[vf.Rubric] | None = None,
+    **env_kwargs: Any,
+) -> vf.SingleTurnEnv:
+    dataset = _load_sharegpt_dataset(data_path)
+
+    client = PolarClient(
+        base_url=f"https://{server_address}",
+        model=reward_model,
+    )
+
+    rubric_cls = reward_scheme or vf.Rubric
+    rubric = rubric_cls(funcs=[polar_reward])
+    rubric.class_objects["reward_client"] = client
+
+    return vf.SingleTurnEnv(dataset=dataset, rubric=rubric, **env_kwargs)
+
+
+class PolarClient:
+    """Minimal async client for POLAR reward model served via vLLM."""
+
+    def __init__(self, *, base_url: str, model: str, timeout: float = 30.0, api_key: str | None = None):
+        self.base_url = base_url.rstrip("/")
+        self.model = model
+        self.timeout = timeout
+        self.api_key = api_key
+
+    async def score(self, payload: list[dict[str, Any]]) -> list[float]:
+        encoded = self._encode(payload)
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.post(
+                f"{self.base_url}/v1/rewards",
+                json={"model": self.model, "input": encoded},
+                headers={"Authorization": f"Bearer {self.api_key}"} if self.api_key else None,
+            )
+        if response.status_code != HTTPStatus.OK:
+            raise RuntimeError(
+                f"POLAR reward request failed: {response.status_code} {response.text}"
+            )
+        data = response.json()
+        return data.get("rewards", [])
+
+    @staticmethod
+    def _encode(payload: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        # Ensure payload matches expected schema; keep implementation simple for now.
+        return payload
+

pyproject.toml

@@ -0,0 +1,21 @@
+[project]
+name = "sharegpt-polar"
+version = "0.1.0"
+description = "blehbklgvggfg"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "datasets>=2.16.0",
+    "httpx>=0.27.0",
+]
+
+[[project.authors]]
+name = "Prime Intellect"
+email = "infra@primeintellect.ai"
+
+[project.entry-points."verifiers.environments"]
+sharegpt-polar = "sharegpt_polar:load_environment"
+
+[tool.setuptools]
+py-modules = ["sharegpt_polar"]
+

sharegpt_polar.egg-info/PKG-INFO

@@ -0,0 +1,67 @@
+Metadata-Version: 2.4
+Name: sharegpt-polar
+Version: 0.1.0
+Summary: ShareGPT reward environment scored by POLAR
+Author-email: Prime Intellect <infra@primeintellect.ai>
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: datasets>=2.16.0
+Requires-Dist: httpx>=0.27.0
+
+# ShareGPT POLAR Environment
+
+This environment scores policy rollouts using the POLAR reward model served via vLLM. It expects a ShareGPT-style JSONL dataset and treats the original assistant response as a reference trajectory.
+
+## Dataset Format
+
+Each line in the dataset must contain a `conversations` list with alternating human/GPT turns:
+
+```json
+{
+  "conversations": [
+    {"from": "human", "value": "Prompt text"},
+    {"from": "gpt", "value": "Reference answer"}
+  ]
+}
+```
+
+The loader extracts the first human message as the prompt and the first GPT message as the reference. Additional turns can be present and are ignored.
+
+## Usage
+
+```python
+import verifiers as vf
+
+env = vf.load_environment(
+    "sharegpt-polar",
+    data_path="/path/to/sharegpt.jsonl",
+    server_address="wealth-intent-submissions-range.trycloudflare.com",
+)
+```
+
+The environment bundles an async reward function which submits `(prompt, reference, output)` to the POLAR reward model using `RewardModelClient` from xtuner.
+
+## Dependencies
+
+- `datasets>=2.16.0`
+- `xtuner[deepspeed]==0.2.0`
+
+Install them in your project via:
+
+```bash
+uv add datasets
+uv add httpx
+```
+
+## Evaluation
+
+Run quick evaluations with any OpenAI-compatible policy model:
+
+```bash
+uv run vf-install sharegpt-polar -p /home/Ubuntu/Mango/verifiers/environments
+uv run vf-eval sharegpt-polar -m gpt-4.1-mini \
+  --env-args '{"data_path": "data/sharegpt.jsonl"}' -n 10 -r 1
+```
+
+Ensure the POLAR reward model is reachable at the supplied `server_address` before evaluation.
+

sharegpt_polar.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+README.md
+pyproject.toml
+sharegpt_polar.py
+sharegpt_polar.egg-info/PKG-INFO
+sharegpt_polar.egg-info/SOURCES.txt
+sharegpt_polar.egg-info/dependency_links.txt
+sharegpt_polar.egg-info/entry_points.txt
+sharegpt_polar.egg-info/requires.txt
+sharegpt_polar.egg-info/top_level.txt

sharegpt_polar.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

sharegpt_polar.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[verifiers.environments]
+sharegpt-polar = sharegpt_polar:load_environment

sharegpt_polar.egg-info/requires.txt

@@ -0,0 +1,2 @@
+datasets>=2.16.0
+httpx>=0.27.0

sharegpt_polar.egg-info/top_level.txt

@@ -0,0 +1 @@
+sharegpt_polar

sharegpt_polar.py

@@ -0,0 +1,129 @@
+"""ShareGPT + POLAR reward environment."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset, load_dataset
+import asyncio
+
+import verifiers as vf
+from verifiers.types import Messages
+from xtuner.utils import RewardModelClient
+
+DEFAULT_MODEL = "internlm/POLAR-7B"
+
+
+def _load_sharegpt_dataset(path: str | Path) -> Dataset:
+    dataset = load_dataset("json", data_files=str(path), split="train")
+
+    def to_single_turn(example: dict[str, Any]) -> dict[str, Any]:
+        human_turn = next(
+            turn["value"] for turn in example["conversations"] if turn["from"] == "human"
+        )
+        assistant_turn = next(
+            turn["value"] for turn in example["conversations"] if turn["from"] == "gpt"
+        )
+        return {
+            "prompt": [{"role": "user", "content": human_turn}],
+            "info": {
+                "reference": [{"role": "assistant", "content": assistant_turn}],
+            },
+        }
+
+    return dataset.map(to_single_turn, remove_columns=dataset.column_names)
+
+
+class PoolingClient:
+    def __init__(
+        self,
+        model_path: str,
+        server_address: str,
+        server_type: str = "lmdeploy",
+        max_length: int = 16384,
+        max_response_length: int = 4096,
+        response_cut_side: str = "left",
+    ):
+        self.client = RewardModelClient(
+            model_path,
+            max_length=max_length,
+            max_response_length=max_response_length,
+            response_cut_side=response_cut_side,
+            server_type=server_type,
+            server_address=server_address,
+        )
+
+    def encode(self, sample: dict[str, Any]) -> str:
+        prompt_text = "\n".join(
+            message["content"] for message in sample.get("prompt", [])
+        )
+        reference_text = "\n".join(
+            message["content"] for message in sample.get("reference", [])
+        )
+        output_text = "\n".join(
+            message["content"] for message in sample.get("output", [])
+        )
+        return f"{prompt_text}\n{reference_text}<|reward|>{prompt_text}\n{output_text}[UNUSED_TOKEN_130]"
+
+    def score(self, payload: list[dict[str, Any]]) -> list[float]:
+        encoded_payload = [self.encode(item) for item in payload]
+        rewards = self.client.lmdeploy_request_reward(encoded_payload)
+        if rewards is None:
+            raise RuntimeError("Failed to get rewards from lmdeploy server")
+        return rewards
+
+
+async def polar_reward(
+    prompt: Messages,
+    completion: Messages,
+    info: dict[str, Any],
+    reward_client: PoolingClient,
+    pooling_semaphore: asyncio.Semaphore,
+    **_: Any,
+) -> float:
+    assistant_turns = [msg for msg in completion if msg.get("role") == "assistant"]
+    if not assistant_turns:
+        return 0.0
+
+    payload = [
+        {
+            "prompt": prompt,
+            "reference": info.get("reference", []),
+            "output": [assistant_turns[-1]],
+        }
+    ]
+    async with pooling_semaphore:
+        loop = asyncio.get_running_loop()
+        rewards = await loop.run_in_executor(None, reward_client.score, payload)
+    if rewards:
+        return float(rewards[-1]) * 10.0
+    raise RuntimeError(f"Unexpected reward response: {rewards}")
+
+
+def load_environment(
+    data_path: str | Path,
+    *,
+    server_address: str,
+    reward_model: str = DEFAULT_MODEL,
+    reward_scheme: type[vf.Rubric] | None = None,
+    server_type: str = "lmdeploy",
+    **env_kwargs: Any,
+) -> vf.SingleTurnEnv:
+    dataset = _load_sharegpt_dataset(data_path)
+
+    client = PoolingClient(
+        model_path=reward_model,
+        server_address=server_address,
+        server_type=server_type,
+    )
+
+    rubric_cls = reward_scheme or vf.Rubric
+    rubric = rubric_cls(funcs=[polar_reward])
+    rubric.class_objects["reward_client"] = client
+    rubric.class_objects.setdefault("pooling_semaphore", asyncio.Semaphore(4))
+
+    return vf.SingleTurnEnv(dataset=dataset, rubric=rubric, **env_kwargs)
+
+
+

train_sharegpt_polar.py

@@ -0,0 +1,108 @@
+"""GRPO training entrypoint for ShareGPT POLAR environment."""
+
+from __future__ import annotations
+
+import os
+
+import verifiers as vf
+from xtuner.utils import RewardModelClient
+
+RUN_NAME = "sharegpt-polar"
+MODEL_NAME = "NewEden/Snwy-SFT-GRPO-base"
+DATA_PATH = "/home/Ubuntu/Mango/verifiers/new.jsonl"
+SERVER_ADDRESS = "greene-cannon-republic-expect.trycloudflare.com"
+REWARD_MODEL = "internlm/POLAR-7B"
+
+# Training hyperparameters
+PER_DEVICE_TRAIN_BATCH_SIZE = 2
+NUM_GENERATIONS = 8
+GRADIENT_ACCUMULATION_STEPS = 2
+LEARNING_RATE = 1e-6
+BETA = 0.1
+MAX_STEPS = 1000
+MAX_GRAD_NORM = 1.0
+NUM_ITERATIONS = 1
+MAX_TOKENS = 512
+TEMPERATURE = 1.0
+TOP_P = 1.0
+SAVE_EVERY_STEPS = 50
+LOGGING_STEPS = 1
+REPORT_TO = ["wandb"]
+LOG_COMPLETIONS = True
+LOG_ON_EACH_NODE = False
+ASYNC_GENERATION_TIMEOUT = 60000
+MAX_CONCURRENT = 1024
+WANDB_PROJECT = "14B-GRPO"
+WANDB_NAME = RUN_NAME
+
+if WANDB_PROJECT:
+    os.environ.setdefault("WANDB_PROJECT", WANDB_PROJECT)
+if WANDB_NAME:
+    os.environ.setdefault("WANDB_NAME", WANDB_NAME)
+
+
+def _check_reward_server() -> None:
+    client = RewardModelClient(
+        REWARD_MODEL,
+        server_type="lmdeploy",
+        server_address=SERVER_ADDRESS,
+    )
+    sanity_samples = [
+        {
+            "prompt": "What is the capital of China?",
+            "reference": "Beijing.",
+            "output": "Beijing.",
+        },
+        {
+            "prompt": "What is the capital of China?",
+            "reference": "Beijing.",
+            "output": "Shanghai.",
+        },
+    ]
+    encoded = client.encode(sanity_samples)
+    rewards = client.lmdeploy_request_reward(encoded)
+    print("[sanity] lmdeploy rewards:", rewards)
+
+
+_check_reward_server()
+
+vf_env = vf.load_environment(
+    env_id="sharegpt-polar",
+    data_path=DATA_PATH,
+    server_address=SERVER_ADDRESS,
+)
+
+model, tokenizer = vf.get_model_and_tokenizer(MODEL_NAME)
+
+training_args = vf.grpo_defaults(run_name=RUN_NAME)
+training_args.per_device_train_batch_size = PER_DEVICE_TRAIN_BATCH_SIZE
+training_args.num_generations = NUM_GENERATIONS
+training_args.gradient_accumulation_steps = GRADIENT_ACCUMULATION_STEPS
+training_args.learning_rate = LEARNING_RATE
+training_args.beta = BETA
+training_args.max_steps = MAX_STEPS
+training_args.max_grad_norm = MAX_GRAD_NORM
+training_args.num_iterations = NUM_ITERATIONS
+training_args.max_tokens = MAX_TOKENS
+training_args.temperature = TEMPERATURE
+training_args.top_p = TOP_P
+training_args.save_strategy = "steps"
+training_args.save_steps = SAVE_EVERY_STEPS
+training_args.logging_steps = LOGGING_STEPS
+training_args.report_to = REPORT_TO
+training_args.log_completions = LOG_COMPLETIONS
+training_args.log_on_each_node = LOG_ON_EACH_NODE
+training_args.async_generation_timeout = ASYNC_GENERATION_TIMEOUT
+training_args.max_concurrent = MAX_CONCURRENT
+
+trainer = vf.GRPOTrainer(
+    env=vf_env,
+    model=model,
+    processing_class=tokenizer,
+    args=training_args,
+    peft_config=vf.lora_defaults(r=128, alpha=64),
+)
+
+if __name__ == "__main__":
+    trainer.train()
+