Add LFS tracking for binary files (PNG, PDF, ipynb)

.gitattributes

@@ -1,5 +1,6 @@
 *.png filter=lfs diff=lfs merge=lfs -text
 *.pdf filter=lfs diff=lfs merge=lfs -text
+*.ipynb filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.jpeg filter=lfs diff=lfs merge=lfs -text
 *.gif filter=lfs diff=lfs merge=lfs -text

train_colab.ipynb

@@ -1,438 +1,3 @@
-{
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": [],
-      "gpuType": "T4"
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
-    "accelerator": "GPU"
-  },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "# 🏛️ SocialContract-v0 — GRPO Training Notebook\n",
-        "\n",
-        "**Environment:** Fiscal Policy Advisory (OpenEnv-compliant)  \n",
-        "**Training Method:** GRPO (Group Relative Policy Optimization) via TRL  \n",
-        "**Efficient Fine-tuning:** Unsloth (4-bit LoRA)  \n",
-        "**Model:** Qwen2.5-1.5B-Instruct  \n",
-        "\n",
-        "This notebook trains an LLM to be an IMF-grade economic policy advisor using\n",
-        "reinforcement learning from the SocialContract-v0 environment rewards."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 1. Install Dependencies"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "%%capture\n",
-        "!pip install unsloth trl datasets transformers accelerate pydantic numpy matplotlib\n",
-        "!pip install --no-deps xformers\n",
-        "print('Dependencies installed!')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 2. Clone the Environment"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Clone the SocialContract-v0 environment repo\n",
-        "!git clone https://huggingface.co/spaces/Tyr-123/SocialContract-v0 /content/socialcontract\n",
-        "import sys\n",
-        "sys.path.insert(0, '/content/socialcontract')\n",
-        "\n",
-        "# Verify environment works\n",
-        "from env.openenv_wrapper import SocialContractOpenEnv\n",
-        "from env.pydantic_models import PolicyAction\n",
-        "from graders.graders import run_grader\n",
-        "\n",
-        "env = SocialContractOpenEnv('task1_stability')\n",
-        "obs = env.reset()\n",
-        "print(f'✅ Environment loaded! GDP={obs.gdp:.1f}, Gini={obs.gini:.3f}')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 3. Quick Baseline Evaluation"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import numpy as np\n",
-        "\n",
-        "# Random baseline\n",
-        "def run_random_episode(task_id, seed=42):\n",
-        "    rng = np.random.default_rng(seed)\n",
-        "    env = SocialContractOpenEnv(task_id, seed=seed)\n",
-        "    env.reset()\n",
-        "    while not env.is_done:\n",
-        "        action = PolicyAction(\n",
-        "            tax_delta=float(rng.uniform(-0.05, 0.05)),\n",
-        "            ubi_delta=float(rng.uniform(-5.0, 5.0)),\n",
-        "            public_good_delta=float(rng.uniform(-0.05, 0.05)),\n",
-        "            interest_rate_delta=float(rng.uniform(-0.02, 0.02)),\n",
-        "            stimulus_package=float(rng.uniform(0, 100)),\n",
-        "            import_tariff_delta=float(rng.uniform(-0.03, 0.03)),\n",
-        "            money_supply_delta=float(rng.uniform(-100, 100)),\n",
-        "            minimum_wage_delta=float(rng.uniform(-1, 1)),\n",
-        "            reasoning='random',\n",
-        "        )\n",
-        "        env.step(action)\n",
-        "    return run_grader(task_id, env._history)\n",
-        "\n",
-        "tasks = ['task1_stability', 'task2_recession', 'task4_stagflation']\n",
-        "print('Random Baseline Scores:')\n",
-        "for task in tasks:\n",
-        "    grade = run_random_episode(task)\n",
-        "    print(f'  {task}: {grade[\"score\"]:.4f} — {grade[\"verdict\"]}')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 4. Load Model with Unsloth"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from unsloth import FastLanguageModel\n",
-        "import torch\n",
-        "\n",
-        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-        "    model_name='unsloth/Qwen2.5-1.5B-Instruct',\n",
-        "    max_seq_length=2048,\n",
-        "    dtype=None,\n",
-        "    load_in_4bit=True,\n",
-        ")\n",
-        "\n",
-        "model = FastLanguageModel.get_peft_model(\n",
-        "    model,\n",
-        "    r=16,\n",
-        "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
-        "                    'gate_proj', 'up_proj', 'down_proj'],\n",
-        "    lora_alpha=32,\n",
-        "    lora_dropout=0.0,\n",
-        "    bias='none',\n",
-        "    use_gradient_checkpointing='unsloth',\n",
-        ")\n",
-        "\n",
-        "if tokenizer.pad_token is None:\n",
-        "    tokenizer.pad_token = tokenizer.eos_token\n",
-        "\n",
-        "print(f'✅ Model loaded: Qwen2.5-1.5B-Instruct with LoRA (r=16)')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 5. Generate Training Prompts"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import random\n",
-        "from datasets import Dataset\n",
-        "\n",
-        "SYSTEM_PROMPT = (\n",
-        "    'You are an expert economic policy advisor. You make precise, data-driven '\n",
-        "    'decisions to optimize economic outcomes across multiple competing objectives. '\n",
-        "    'You must respond ONLY with valid JSON matching the exact format requested. '\n",
-        "    'No markdown, no extra text — only the JSON object.'\n",
-        ")\n",
-        "\n",
-        "def generate_prompts(n=80):\n",
-        "    prompts = []\n",
-        "    for task_id in tasks:\n",
-        "        for seed in range(1, n // len(tasks) + 2):\n",
-        "            if len(prompts) >= n:\n",
-        "                break\n",
-        "            env = SocialContractOpenEnv(task_id, seed=seed)\n",
-        "            obs = env.reset()\n",
-        "            warmup = random.randint(0, env.task_cfg['max_steps'] // 3)\n",
-        "            rng = np.random.default_rng(seed)\n",
-        "            for _ in range(warmup):\n",
-        "                action = PolicyAction(\n",
-        "                    tax_delta=float(rng.uniform(-0.02, 0.02)),\n",
-        "                    ubi_delta=float(rng.uniform(-1, 1)),\n",
-        "                    public_good_delta=float(rng.uniform(-0.02, 0.02)),\n",
-        "                    interest_rate_delta=float(rng.uniform(-0.01, 0.01)),\n",
-        "                    stimulus_package=float(rng.uniform(0, 30)),\n",
-        "                    import_tariff_delta=0.0,\n",
-        "                    money_supply_delta=float(rng.uniform(-30, 30)),\n",
-        "                    minimum_wage_delta=0.0,\n",
-        "                    reasoning='warmup',\n",
-        "                )\n",
-        "                obs, _, done, _ = env.step(action)\n",
-        "                if done: break\n",
-        "            if not env.is_done:\n",
-        "                prompts.append({\n",
-        "                    'messages': [\n",
-        "                        {'role': 'system', 'content': SYSTEM_PROMPT},\n",
-        "                        {'role': 'user', 'content': obs.to_prompt()},\n",
-        "                    ],\n",
-        "                    'task_id': task_id,\n",
-        "                    'seed': seed,\n",
-        "                    'warmup_steps': warmup,\n",
-        "                })\n",
-        "    random.shuffle(prompts)\n",
-        "    return prompts[:n]\n",
-        "\n",
-        "prompt_data = generate_prompts(80)\n",
-        "dataset = Dataset.from_list(prompt_data)\n",
-        "print(f'✅ Generated {len(dataset)} training prompts')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 6. Define Reward Function"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import re, json\n",
-        "\n",
-        "def parse_action(text):\n",
-        "    try:\n",
-        "        if '```' in text:\n",
-        "            match = re.search(r'```(?:json)?\\s*(.*?)```', text, re.DOTALL)\n",
-        "            if match: text = match.group(1)\n",
-        "        match = re.search(r'\\{[^{}]*\\}', text, re.DOTALL)\n",
-        "        if match: text = match.group(0)\n",
-        "        parsed = json.loads(text)\n",
-        "        return PolicyAction(\n",
-        "            tax_delta=float(np.clip(parsed.get('tax_delta', 0), -0.1, 0.1)),\n",
-        "            ubi_delta=float(np.clip(parsed.get('ubi_delta', 0), -10, 10)),\n",
-        "            public_good_delta=float(np.clip(parsed.get('public_good_delta', 0), -0.1, 0.1)),\n",
-        "            interest_rate_delta=float(np.clip(parsed.get('interest_rate_delta', 0), -0.03, 0.03)),\n",
-        "            stimulus_package=float(np.clip(parsed.get('stimulus_package', 0), 0, 500)),\n",
-        "            import_tariff_delta=float(np.clip(parsed.get('import_tariff_delta', 0), -0.05, 0.05)),\n",
-        "            money_supply_delta=float(np.clip(parsed.get('money_supply_delta', 0), -500, 500)),\n",
-        "            minimum_wage_delta=float(np.clip(parsed.get('minimum_wage_delta', 0), -2, 2)),\n",
-        "            reasoning=str(parsed.get('reasoning', ''))[:200],\n",
-        "        )\n",
-        "    except:\n",
-        "        return None\n",
-        "\n",
-        "def env_reward(completion, task_id, seed, warmup_steps):\n",
-        "    action = parse_action(completion)\n",
-        "    if action is None:\n",
-        "        return 0.0\n",
-        "    env = SocialContractOpenEnv(task_id, seed=seed)\n",
-        "    obs = env.reset()\n",
-        "    rng = np.random.default_rng(seed)\n",
-        "    for _ in range(warmup_steps):\n",
-        "        wa = PolicyAction(\n",
-        "            tax_delta=float(rng.uniform(-0.02, 0.02)),\n",
-        "            ubi_delta=float(rng.uniform(-1, 1)),\n",
-        "            public_good_delta=float(rng.uniform(-0.02, 0.02)),\n",
-        "            interest_rate_delta=float(rng.uniform(-0.01, 0.01)),\n",
-        "            stimulus_package=float(rng.uniform(0, 30)),\n",
-        "            import_tariff_delta=0.0,\n",
-        "            money_supply_delta=float(rng.uniform(-30, 30)),\n",
-        "            minimum_wage_delta=0.0,\n",
-        "            reasoning='warmup',\n",
-        "        )\n",
-        "        obs, _, done, _ = env.step(wa)\n",
-        "        if done: return 0.0\n",
-        "    obs, reward, done, _ = env.step(action)\n",
-        "    # Run remaining with a follow-through\n",
-        "    while not env.is_done:\n",
-        "        heur = PolicyAction(\n",
-        "            tax_delta=action.tax_delta * 0.5,\n",
-        "            ubi_delta=action.ubi_delta * 0.3,\n",
-        "            public_good_delta=action.public_good_delta * 0.5,\n",
-        "            interest_rate_delta=action.interest_rate_delta * 0.3,\n",
-        "            stimulus_package=0.0, import_tariff_delta=0.0,\n",
-        "            money_supply_delta=action.money_supply_delta * 0.2,\n",
-        "            minimum_wage_delta=action.minimum_wage_delta * 0.3,\n",
-        "            reasoning='followthrough',\n",
-        "        )\n",
-        "        obs, _, done, _ = env.step(heur)\n",
-        "    return run_grader(task_id, env._history)['score']\n",
-        "\n",
-        "def reward_function(completions, **kwargs):\n",
-        "    rewards = []\n",
-        "    task_ids = kwargs.get('task_id', ['task1_stability'] * len(completions))\n",
-        "    seeds = kwargs.get('seed', [42] * len(completions))\n",
-        "    warmups = kwargs.get('warmup_steps', [0] * len(completions))\n",
-        "    for i, c in enumerate(completions):\n",
-        "        text = c[-1].get('content', '') if isinstance(c, list) else str(c)\n",
-        "        r = env_reward(text,\n",
-        "                       task_ids[i] if i < len(task_ids) else 'task1_stability',\n",
-        "                       seeds[i] if i < len(seeds) else 42,\n",
-        "                       warmups[i] if i < len(warmups) else 0)\n",
-        "        rewards.append(r)\n",
-        "    return rewards\n",
-        "\n",
-        "print('✅ Reward function defined')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 7. Train with GRPO"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from trl import GRPOTrainer, GRPOConfig\n",
-        "\n",
-        "training_args = GRPOConfig(\n",
-        "    output_dir='./socialcontract-grpo',\n",
-        "    num_train_epochs=2,\n",
-        "    per_device_train_batch_size=2,\n",
-        "    gradient_accumulation_steps=4,\n",
-        "    learning_rate=5e-5,\n",
-        "    max_completion_length=256,\n",
-        "    num_generations=4,\n",
-        "    logging_steps=5,\n",
-        "    save_steps=50,\n",
-        "    max_steps=150,\n",
-        "    report_to='none',\n",
-        "    bf16=torch.cuda.is_available(),\n",
-        "    temperature=0.7,\n",
-        ")\n",
-        "\n",
-        "trainer = GRPOTrainer(\n",
-        "    model=model,\n",
-        "    args=training_args,\n",
-        "    train_dataset=dataset,\n",
-        "    processing_class=tokenizer,\n",
-        "    reward_funcs=[reward_function],\n",
-        ")\n",
-        "\n",
-        "print('🚀 Starting GRPO training...')\n",
-        "result = trainer.train()\n",
-        "print(f'✅ Training complete! Steps: {result.global_step}')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 8. Post-Training Evaluation"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "def evaluate_trained_model(task_id, seed=42):\n",
-        "    env = SocialContractOpenEnv(task_id, seed=seed)\n",
-        "    obs = env.reset()\n",
-        "    while not env.is_done:\n",
-        "        prompt = obs.to_prompt()\n",
-        "        messages = [\n",
-        "            {'role': 'system', 'content': SYSTEM_PROMPT},\n",
-        "            {'role': 'user', 'content': prompt},\n",
-        "        ]\n",
-        "        inputs = tokenizer.apply_chat_template(\n",
-        "            messages, return_tensors='pt', add_generation_prompt=True\n",
-        "        ).to(model.device)\n",
-        "        with torch.no_grad():\n",
-        "            outputs = model.generate(\n",
-        "                inputs, max_new_tokens=256, temperature=0.1,\n",
-        "                do_sample=True, pad_token_id=tokenizer.eos_token_id,\n",
-        "            )\n",
-        "        resp = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
-        "        action = parse_action(resp)\n",
-        "        if action is None:\n",
-        "            action = PolicyAction(\n",
-        "                tax_delta=0.0, ubi_delta=0.0, public_good_delta=0.01,\n",
-        "                interest_rate_delta=0.0, stimulus_package=0.0,\n",
-        "                import_tariff_delta=0.0, money_supply_delta=0.0,\n",
-        "                minimum_wage_delta=0.0, reasoning='parse failure',\n",
-        "            )\n",
-        "        obs, _, _, _ = env.step(action)\n",
-        "    return run_grader(task_id, env._history)\n",
-        "\n",
-        "print('Post-Training Scores:')\n",
-        "for task in tasks:\n",
-        "    grade = evaluate_trained_model(task)\n",
-        "    print(f'  {task}: {grade[\"score\"]:.4f} — {grade[\"verdict\"]}')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 9. Save & Push to HuggingFace Hub"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Save locally\n",
-        "model.save_pretrained('./socialcontract-grpo-final')\n",
-        "tokenizer.save_pretrained('./socialcontract-grpo-final')\n",
-        "\n",
-        "# Optionally push to HuggingFace Hub\n",
-        "# from huggingface_hub import login\n",
-        "# login(token='your-hf-token')\n",
-        "# model.push_to_hub('your-username/socialcontract-policy-v0')\n",
-        "# tokenizer.push_to_hub('your-username/socialcontract-policy-v0')\n",
-        "\n",
-        "print('✅ Model saved!')"
-      ]
-    }
-  ]
-}
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5d1cb03f0ec786177271c5863cde41bb9c0c693b9154258eb9619830f0d1315
+size 17145