{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "title",
   "metadata": {},
   "source": [
    "# Slide Skill Optimizer \u2014 Training Notebook\n",
    "\n",
    "Fine-tunes a small open model (Qwen2.5-7B) to replace Claude Opus 4.6 as the\n",
    "**optimizer** in the Slide Skill OpenEnv loop.\n",
    "\n",
    "**Pipeline:**\n",
    "1. **Data collection** \u2014 Run N episodes against the live OpenEnv server using\n",
    "   Claude Opus 4.6 as the oracle optimizer. Record every (prompt, DESIGN_RULES rewrite, reward) tuple.\n",
    "2. **Filtering** \u2014 Keep only steps where `reward > 0` (the rewrite improved the score).\n",
    "3. **SFT training** \u2014 Fine-tune Qwen2.5-7B-Instruct with Unsloth + TRL SFTTrainer.\n",
    "4. **Save / push** \u2014 Export QLoRA weights; optionally push to HuggingFace Hub.\n",
    "\n",
    "**Models in the loop:**\n",
    "- Generator: Claude Sonnet 4.6 (writes pptxgenjs JS)\n",
    "- Evaluator: Gemini 3.1 Pro (scores the slide with vision)\n",
    "- Oracle optimizer (data collection): Claude Opus 4.6\n",
    "- Trained optimizer (inference): Qwen2.5-7B fine-tuned here\n",
    "\n",
    "**Runtime:** T4 GPU (free tier) is sufficient. A100 is faster.\n",
    "**Estimated collection time:** ~90s per step \u00d7 7 steps \u00d7 N episodes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "install-header",
   "metadata": {},
   "source": [
    "## 1. Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unsloth must be installed before transformers to patch correctly.\n",
    "!pip install -q \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
    "!pip install -q anthropic httpx trl datasets transformers accelerate bitsandbytes\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "config-header",
   "metadata": {},
   "source": [
    "## 2. Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "config",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# API keys\n",
    "# ------------------------------------------------------------------\n",
    "# Set these here or via Colab Secrets (Secrets tab in left sidebar)\n",
    "os.environ.setdefault('ANTHROPIC_API_KEY', '')  # for oracle optimizer (Opus 4.6)\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# OpenEnv server URL\n",
    "# ------------------------------------------------------------------\n",
    "# Option A: HuggingFace Spaces deployment\n",
    "#   SERVER_URL = 'https://your-username-slide-skill-openenv.hf.space'\n",
    "# Option B: Local server exposed via ngrok\n",
    "#   !pip install -q pyngrok\n",
    "#   from pyngrok import ngrok; ngrok.set_auth_token('YOUR_TOKEN')\n",
    "#   tunnel = ngrok.connect(8000); SERVER_URL = tunnel.public_url\n",
    "SERVER_URL = 'https://your-hf-space-url.hf.space'  # <-- EDIT THIS\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# Data collection\n",
    "# ------------------------------------------------------------------\n",
    "N_EPISODES = 20          # Number of full optimization episodes to collect\n",
    "MAX_STEPS_PER_EPISODE = 7\n",
    "MIN_REWARD = 0.0         # Only train on steps where reward > this threshold\n",
    "DATA_FILE = 'trajectories.json'\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# Training\n",
    "# ------------------------------------------------------------------\n",
    "BASE_MODEL = 'unsloth/Qwen2.5-7B-Instruct-bnb-4bit'\n",
    "MAX_SEQ_LENGTH = 4096\n",
    "OUTPUT_DIR = './optimizer-model'\n",
    "HF_REPO = ''  # Optional: 'your-username/slide-optimizer' to push to Hub\n",
    "\n",
    "print(f'Server: {SERVER_URL}')\n",
    "print(f'Collecting {N_EPISODES} episodes, training on steps with reward > {MIN_REWARD}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "collect-header",
   "metadata": {},
   "source": [
    "## 3. Data Collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import textwrap\n",
    "import time\n",
    "from dataclasses import dataclass, asdict\n",
    "from typing import Any\n",
    "\n",
    "import anthropic\n",
    "import httpx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "oracle",
   "metadata": {},
   "outputs": [],
   "source": [
    "ORACLE_MODEL = 'claude-opus-4-6'\n",
    "OPTIMIZER_SYSTEM = (\n",
    "    'You are a McKinsey slide design optimizer. '\n",
    "    'You improve a PowerPoint generation skill by rewriting its DESIGN_RULES.md file. '\n",
    "    'Output ONLY the markdown file content \u2014 no explanation, no code fences.'\n",
    ")\n",
    "\n",
    "\n",
    "def build_optimizer_prompt(obs: dict) -> str:\n",
    "    \"\"\"Format an observation dict into the optimizer input prompt.\"\"\"\n",
    "    scores = obs['scores']\n",
    "    strengths = '\\n'.join(f'- {s}' for s in obs.get('strengths', []))\n",
    "    weaknesses = '\\n'.join(f'- {w}' for w in obs.get('weaknesses', []))\n",
    "    return textwrap.dedent(f\"\"\"\\\n",
    "        ## Current Score: {obs['total']}/100\n",
    "\n",
    "        ## Score Breakdown\n",
    "        - background_layout: {scores['background_layout']}/15\n",
    "        - color_palette: {scores['color_palette']}/15\n",
    "        - typography: {scores['typography']}/15\n",
    "        - title_quality: {scores['title_quality']}/15\n",
    "        - data_presentation: {scores['data_presentation']}/15\n",
    "        - structural_elements: {scores['structural_elements']}/15\n",
    "        - overall_impression: {scores['overall_impression']}/10\n",
    "\n",
    "        ## Evaluator Feedback\n",
    "        Strengths:\n",
    "        {strengths}\n",
    "\n",
    "        Weaknesses:\n",
    "        {weaknesses}\n",
    "\n",
    "        Verdict: {obs['one_line_verdict']}\n",
    "\n",
    "        ## Current DESIGN_RULES.md\n",
    "        {obs['design_rules_content']}\n",
    "\n",
    "        ## Current EXAMPLES.md\n",
    "        {obs['examples_content']}\n",
    "\n",
    "        Write an improved DESIGN_RULES.md that addresses the weaknesses above \\\n",
    "        while preserving what works well. Focus on the lowest-scoring dimensions.\n",
    "    \"\"\")\n",
    "\n",
    "\n",
    "def call_oracle(prompt: str, anthropic_client: anthropic.Anthropic) -> str:\n",
    "    \"\"\"Call Claude Opus 4.6 to generate an improved DESIGN_RULES.md.\"\"\"\n",
    "    response = anthropic_client.messages.create(\n",
    "        model=ORACLE_MODEL,\n",
    "        max_tokens=4096,\n",
    "        system=OPTIMIZER_SYSTEM,\n",
    "        messages=[{'role': 'user', 'content': prompt}],\n",
    "    )\n",
    "    return response.content[0].text.strip()\n",
    "\n",
    "\n",
    "print('Oracle functions defined.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collect-fn",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\nclass TrainingStep:\n    prompt: str       # optimizer input (observation formatted as text)\n    completion: str   # oracle output (new DESIGN_RULES.md)\n    reward: float     # environment reward for this step\n    total: int        # slide score after this step\n    episode: int\n    step: int\n\n\nBASELINE_EXAMPLES = '(Empty \u2014 no prior optimization rounds)\\n'\n\n\ndef collect_episode(\n    episode_idx: int,\n    http: httpx.Client,\n    anthropic_client: anthropic.Anthropic,\n    max_steps: int = MAX_STEPS_PER_EPISODE,\n) -> list[TrainingStep]:\n    \"\"\"Run one full optimization episode and return training steps.\"\"\"\n    steps = []\n\n    # Reset session.\n    resp = http.post(f'{SERVER_URL}/reset', json={})\n    resp.raise_for_status()\n    session_id = resp.json()['session_id']\n    print(f'  Episode {episode_idx}: session={session_id}')\n\n    # Baseline step (no optimizer action \u2014 just trigger generation).\n    resp = http.post(f'{SERVER_URL}/step', json={\n        'session_id': session_id,\n        'action': {\n            'action_type': 'replace_file',\n            'file': 'EXAMPLES.md',\n            'new_content': BASELINE_EXAMPLES,\n        },\n    })\n    resp.raise_for_status()\n    obs = resp.json()\n    print(f'    baseline score={obs[\"total\"]}/100')\n\n    try:\n        # Optimization steps.\n        for step_idx in range(1, max_steps + 1):\n            if obs.get('done'):\n                break\n\n            # Build prompt and call oracle.\n            prompt = build_optimizer_prompt(obs)\n            completion = call_oracle(prompt, anthropic_client)\n\n            # Submit to environment.\n            resp = http.post(f'{SERVER_URL}/step', json={\n                'session_id': session_id,\n                'action': {\n                    'action_type': 'replace_file',\n                    'file': 'DESIGN_RULES.md',\n                    'new_content': completion,\n                },\n            })\n            resp.raise_for_status()\n            obs = resp.json()\n\n            steps.append(TrainingStep(\n                prompt=prompt,\n                completion=completion,\n                reward=obs['reward'],\n                total=obs['total'],\n                episode=episode_idx,\n                step=step_idx,\n            ))\n            delta = f\"{obs['reward']*100:+.0f}pts\"\n            print(f'    step {step_idx}: score={obs[\"total\"]}/100 ({delta})')\n    finally:\n        # Always clean up session, even on error.\n        http.delete(f'{SERVER_URL}/sessions/{session_id}')\n    return steps\n\n\nprint('Collection function defined.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "run-collection",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_steps: list[TrainingStep] = []\n",
    "\n",
    "anthropic_client = anthropic.Anthropic(api_key=os.environ['ANTHROPIC_API_KEY'])\n",
    "\n",
    "# Long timeout \u2014 each step takes 60-120s (LLM + Node + LibreOffice + Gemini).\n",
    "with httpx.Client(timeout=300.0) as http:\n",
    "    for ep in range(N_EPISODES):\n",
    "        print(f'\\nEpisode {ep + 1}/{N_EPISODES}')\n",
    "        try:\n",
    "            episode_steps = collect_episode(ep + 1, http, anthropic_client)\n",
    "            all_steps.extend(episode_steps)\n",
    "        except Exception as e:\n",
    "            print(f'  Episode {ep + 1} failed: {e} \u2014 skipping')\n",
    "            continue\n",
    "\n",
    "        # Save incrementally so a Colab crash doesn't lose all data.\n",
    "        with open(DATA_FILE, 'w') as f:\n",
    "            json.dump([asdict(s) for s in all_steps], f, indent=2)\n",
    "        print(f'  Saved {len(all_steps)} steps to {DATA_FILE}')\n",
    "\n",
    "print(f'\\nCollection complete: {len(all_steps)} total steps across {N_EPISODES} episodes.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inspect",
   "metadata": {},
   "outputs": [],
   "source": [
    "import statistics\n\n# Load from disk (in case of kernel restart).\nwith open(DATA_FILE) as f:\n    raw = json.load(f)\nall_steps = [TrainingStep(**s) for s in raw]\n\nrewards = [s.reward for s in all_steps]\npositive = [s for s in all_steps if s.reward > MIN_REWARD]\n\nprint(f'Total steps collected : {len(all_steps)}')\nprint(f'Mean reward           : {statistics.mean(rewards):.3f}' if rewards else 'Mean reward           : n/a')\nprint(f'Positive-reward steps : {len(positive)} ({100*len(positive)/max(len(all_steps),1):.0f}%)')\nprint(f'Training examples     : {len(positive)}')\n\nif len(positive) < 10:\n    print('\\n\u26a0 Very few positive examples. Consider collecting more episodes or lowering MIN_REWARD.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "train-header",
   "metadata": {},
   "source": [
    "## 4. Fine-tuning with Unsloth + TRL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "format-dataset",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "from unsloth import FastLanguageModel\n",
    "\n",
    "# Load model first so we can use its tokenizer for chat formatting.\n",
    "print(f'Loading {BASE_MODEL} ...')\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=BASE_MODEL,\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    load_in_4bit=True,\n",
    ")\n",
    "print('Model loaded.')\n",
    "\n",
    "\n",
    "def format_example(step: TrainingStep) -> dict:\n",
    "    \"\"\"Format a training step as a chat-template string for SFT.\"\"\"\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        [\n",
    "            {'role': 'system', 'content': OPTIMIZER_SYSTEM},\n",
    "            {'role': 'user', 'content': step.prompt},\n",
    "            {'role': 'assistant', 'content': step.completion},\n",
    "        ],\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=False,\n",
    "    )\n",
    "    return {'text': text}\n",
    "\n",
    "\n",
    "formatted = [format_example(s) for s in positive]\n",
    "dataset = Dataset.from_list(formatted).train_test_split(test_size=0.1, seed=42)\n",
    "\n",
    "print(f'Train examples: {len(dataset[\"train\"])}')\n",
    "print(f'Eval  examples: {len(dataset[\"test\"])}')\n",
    "print('\\nSample (truncated):')\n",
    "print(dataset['train'][0]['text'][:500], '...')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lora",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add LoRA adapters via Unsloth.\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r=16,\n",
    "    lora_alpha=16,\n",
    "    target_modules=[\n",
    "        'q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
    "        'gate_proj', 'up_proj', 'down_proj',\n",
    "    ],\n",
    "    lora_dropout=0,\n",
    "    bias='none',\n",
    "    use_gradient_checkpointing='unsloth',  # Unsloth's memory-efficient checkpointing\n",
    "    random_state=42,\n",
    ")\n",
    "print('LoRA adapters attached.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "train",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTTrainer, SFTConfig\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    train_dataset=dataset['train'],\n",
    "    eval_dataset=dataset['test'],\n",
    "    args=SFTConfig(\n",
    "        output_dir=OUTPUT_DIR,\n",
    "        num_train_epochs=3,\n",
    "        per_device_train_batch_size=1,\n",
    "        per_device_eval_batch_size=1,\n",
    "        gradient_accumulation_steps=4,   # effective batch size = 4\n",
    "        learning_rate=2e-4,\n",
    "        lr_scheduler_type='cosine',\n",
    "        warmup_ratio=0.05,\n",
    "        fp16=not FastLanguageModel.is_bfloat16_supported(),\n",
    "        bf16=FastLanguageModel.is_bfloat16_supported(),\n",
    "        logging_steps=5,\n",
    "        eval_strategy='epoch',\n",
    "        save_strategy='epoch',\n",
    "        load_best_model_at_end=True,\n",
    "        report_to='none',  # set to 'wandb' if you want W&B logging\n",
    "        dataset_text_field='text',\n",
    "        max_seq_length=MAX_SEQ_LENGTH,\n",
    "        packing=True,    # Packs short examples together to fill context \u2014 improves throughput\n",
    "    ),\n",
    ")\n",
    "\n",
    "print('Starting training...')\n",
    "trainer.train()\n",
    "print('Training complete.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "save-header",
   "metadata": {},
   "source": [
    "## 5. Save & Push"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "save",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save LoRA adapters (small \u2014 only the delta weights).\n",
    "model.save_pretrained(OUTPUT_DIR)\n",
    "tokenizer.save_pretrained(OUTPUT_DIR)\n",
    "print(f'Saved to {OUTPUT_DIR}/')\n",
    "\n",
    "# Optional: merge + save full model (larger, but self-contained for deployment).\n",
    "# model.save_pretrained_merged(OUTPUT_DIR + '-merged', tokenizer, save_method='merged_16bit')\n",
    "\n",
    "# Optional: push to HuggingFace Hub.\n",
    "if HF_REPO:\n",
    "    model.push_to_hub(HF_REPO, token=os.environ.get('HF_TOKEN', ''))\n",
    "    tokenizer.push_to_hub(HF_REPO, token=os.environ.get('HF_TOKEN', ''))\n",
    "    print(f'Pushed to https://huggingface.co/{HF_REPO}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "inference-header",
   "metadata": {},
   "source": [
    "## 6. Smoke Test \u2014 Run Trained Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inference",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Switch model to inference mode (disables dropout, fuses LoRA for speed).\nFastLanguageModel.for_inference(model)\n\n# Use the last collected observation as a test prompt.\nif not positive:\n    raise RuntimeError('No positive-reward steps to test with \u2014 run data collection first.')\ntest_step = positive[-1]\ntest_prompt = test_step.prompt\n\ninputs = tokenizer.apply_chat_template(\n    [\n        {'role': 'system', 'content': OPTIMIZER_SYSTEM},\n        {'role': 'user', 'content': test_prompt},\n    ],\n    tokenize=True,\n    add_generation_prompt=True,\n    return_tensors='pt',\n).to(model.device)\n\noutputs = model.generate(\n    input_ids=inputs,\n    max_new_tokens=1024,\n    temperature=0.7,\n    do_sample=True,\n)\n\ngenerated = tokenizer.decode(\n    outputs[0][inputs.shape[1]:],\n    skip_special_tokens=True,\n)\n\nprint('=== Trained model output (first 1000 chars) ===')\nprint(generated[:1000])"
   ]
  }
 ]
}