Upload training/make_notebook.py with huggingface_hub

training/make_notebook.py

@@ -0,0 +1,162 @@
+import json
+
+nb = {
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+   "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
+   "language_info": {"name": "python", "version": "3.10.0"}
+ },
+ "cells": [
+  {
+    "cell_type": "markdown", "id": "a1", "metadata": {},
+    "source": [
+      "# DataCentric-Env — GRPO Training Notebook\n\n",
+      "Trains Qwen2.5-3B-Instruct as a data quality agent using GRPO.\n\n",
+      "**Sections:**\n",
+      "1. Install dependencies\n",
+      "2. Model setup (Qwen2.5-3B-Instruct, 4-bit LoRA)\n",
+      "3. Rollout function\n",
+      "4. Collect training data\n",
+      "5. GRPO training loop\n",
+      "6. Save model via Unsloth merge path"
+    ]
+  },
+  {
+    "cell_type": "code", "id": "c1", "metadata": {}, "outputs": [], "execution_count": None,
+    "source": [
+      "# Cell 1: Install dependencies\n",
+      "!pip install unsloth trl transformers accelerate peft datasets requests"
+    ]
+  },
+  {
+    "cell_type": "code", "id": "c2", "metadata": {}, "outputs": [], "execution_count": None,
+    "source": [
+      "# Cell 2: Imports and config\n",
+      "from unsloth import FastLanguageModel\n",
+      "from trl import GRPOTrainer, GRPOConfig\n",
+      "from datasets import Dataset\n",
+      "import requests, json, torch\n",
+      "\n",
+      "ENV_URL = 'https://your-hf-username-datacentric-env.hf.space'  # set your HF Space URL\n",
+      "\n",
+      "SYSTEM_PROMPT = (\n",
+      "    'You are a data quality agent. You receive dataset statistics and must choose '\n",
+      "    'which specialist tool to call to improve the dataset so a downstream classifier '\n",
+      "    'performs better.\\n\\n'\n",
+      "    'Always respond with valid JSON in this exact format:\\n'\n",
+      "    '{\"agent\": \"<tool_name>\", \"target\": \"<column_or_all>\", \"strategy\": \"<strategy_name>\"}\\n\\n'\n",
+      "    'Available tools: cleaner, augmenter, balancer, relabeler, validator\\n'\n",
+      "    'Cleaner strategies: median_impute, mean_impute, drop_rows\\n'\n",
+      "    'Balancer strategies: undersample\\n'\n",
+      "    'Relabeler: use when labels are noisy, costs 2 budget points.'\n",
+      ")\n",
+      "print('Imports OK')"
+    ]
+  },
+  {
+    "cell_type": "code", "id": "c3", "metadata": {}, "outputs": [], "execution_count": None,
+    "source": [
+      "# Cell 3: Model setup\n",
+      "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+      "    model_name='unsloth/Qwen2.5-3B-Instruct',\n",
+      "    max_seq_length=1024,\n",
+      "    load_in_4bit=True,\n",
+      ")\n",
+      "model = FastLanguageModel.get_peft_model(model, r=16, lora_alpha=32)\n",
+      "print('Model loaded')"
+    ]
+  },
+  {
+    "cell_type": "code", "id": "c4", "metadata": {}, "outputs": [], "execution_count": None,
+    "source": [
+      "# Cell 4: Rollout function\n",
+      "def build_prompt(obs):\n",
+      "    return SYSTEM_PROMPT + '\\n\\nCurrent state:\\n' + json.dumps(obs, indent=2) + '\\n\\nYour action:'\n",
+      "\n",
+      "def rollout(prompt='start'):\n",
+      "    obs = requests.post(ENV_URL + '/reset').json()\n",
+      "    trajectories = []\n",
+      "    for step in range(10):\n",
+      "        full_prompt = build_prompt(obs)\n",
+      "        inputs = tokenizer(full_prompt, return_tensors='pt').to('cuda')\n",
+      "        with torch.no_grad():\n",
+      "            outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.7)\n",
+      "        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
+      "        try:\n",
+      "            action = json.loads(response.strip())\n",
+      "        except Exception:\n",
+      "            action = {'agent': 'validator'}\n",
+      "        result = requests.post(ENV_URL + '/step', json=action).json()\n",
+      "        reward = result.get('reward', -1.0)\n",
+      "        trajectories.append({'prompt': full_prompt, 'response': response, 'reward': reward})\n",
+      "        obs = result.get('observation', obs)\n",
+      "        if result.get('done'):\n",
+      "            break\n",
+      "    return trajectories\n",
+      "\n",
+      "print('Rollout function defined')"
+    ]
+  },
+  {
+    "cell_type": "code", "id": "c5", "metadata": {}, "outputs": [], "execution_count": None,
+    "source": [
+      "# Cell 5: Collect rollouts and build dataset\n",
+      "print('Collecting rollouts...')\n",
+      "all_trajectories = []\n",
+      "for episode in range(50):\n",
+      "    all_trajectories.extend(rollout('start'))\n",
+      "    if episode % 10 == 0:\n",
+      "        print(f'  Episode {episode}/50 collected')\n",
+      "\n",
+      "dataset = Dataset.from_list([\n",
+      "    {'prompt': t['prompt'], 'chosen': t['response'], 'reward': t['reward']}\n",
+      "    for t in all_trajectories\n",
+      "])\n",
+      "print(f'Dataset size: {len(dataset)}')"
+    ]
+  },
+  {
+    "cell_type": "code", "id": "c6", "metadata": {}, "outputs": [], "execution_count": None,
+    "source": [
+      "# Cell 6: GRPO training\n",
+      "config = GRPOConfig(\n",
+      "    output_dir='./datacentric-grpo',\n",
+      "    num_train_epochs=3,\n",
+      "    per_device_train_batch_size=4,\n",
+      "    learning_rate=5e-5,\n",
+      "    logging_steps=10,\n",
+      "    save_steps=100,\n",
+      "    report_to='none',\n",
+      ")\n",
+      "\n",
+      "trainer = GRPOTrainer(\n",
+      "    model=model,\n",
+      "    args=config,\n",
+      "    train_dataset=dataset,\n",
+      "    tokenizer=tokenizer,\n",
+      ")\n",
+      "\n",
+      "trainer.train()"
+    ]
+  },
+  {
+    "cell_type": "code", "id": "c7", "metadata": {}, "outputs": [], "execution_count": None,
+    "source": [
+      "# Cell 7: Save via Unsloth merge path\n",
+      "# IMPORTANT: do NOT use naive save_pretrained — use Unsloth merge path\n",
+      "model.save_pretrained_merged(\n",
+      "    'datacentric-grpo-final',\n",
+      "    tokenizer,\n",
+      "    save_method='merged_16bit',\n",
+      ")\n",
+      "print('Training complete. Model saved to datacentric-grpo-final/')"
+    ]
+  }
+ ]
+}
+
+with open("training/train.ipynb", "w") as f:
+    json.dump(nb, f, indent=1)
+
+print("Notebook created successfully.")