{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpsGuard SFT — Colab T4 (or any GPU)\n",
    "\n",
    "Trains Qwen 7B 4-bit + LoRA on 611 OpsGuard rule-based traces. Pushes to `sai1906/opsguard-sft`.\n",
    "\n",
    "**Steps:** Runtime → Change runtime type → T4 GPU. Then Run All. ~30-45 min on T4. Free.\n",
    "\n",
    "**You need:** HF_TOKEN (paste in Cell 2)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -U transformers>=4.46 trl>=0.18 peft>=0.13 bitsandbytes>=0.44 datasets accelerate>=1.0 huggingface_hub matplotlib networkx>=3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Auth + clone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import os, getpass\nos.chdir('/content')\nif not os.environ.get('HF_TOKEN'):\n    os.environ['HF_TOKEN'] = getpass.getpass('Paste HF token (write scope): ')\nfrom huggingface_hub import login\nlogin(token=os.environ['HF_TOKEN'])\n\n!rm -rf /content/opsguard\n!git clone https://huggingface.co/spaces/sai1906/opsguard /content/opsguard\nimport sys; sys.path.insert(0, '/content/opsguard')\nos.chdir('/content/opsguard')\n!ls data/"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load model (Qwen 7B 4-bit) + LoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
    "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "\n",
    "MODEL = 'Qwen/Qwen2.5-7B-Instruct'\n",
    "bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "                          bnb_4bit_use_double_quant=True, bnb_4bit_quant_type='nf4')\n",
    "tok = AutoTokenizer.from_pretrained(MODEL)\n",
    "if tok.pad_token_id is None:\n",
    "    tok.pad_token_id = tok.eos_token_id\n",
    "model = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=bnb,\n",
    "                                              torch_dtype=torch.bfloat16, device_map='auto')\n",
    "model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)\n",
    "lc = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0, bias='none',\n",
    "                target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],\n",
    "                task_type='CAUSAL_LM')\n",
    "model = get_peft_model(model, lc)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Load 611 SFT traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from datasets import Dataset\n",
    "rows = [json.loads(l) for l in open('data/sft_traces.jsonl')]\n",
    "print(f'loaded {len(rows)} traces')\n",
    "texts = [r['prompt'] + '\\n\\nACTION:\\n' + r['completion'] + tok.eos_token for r in rows]\n",
    "ds = Dataset.from_list([{'text': t} for t in texts])\n",
    "print('first sample (truncated):', texts[0][:300], '...')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. SFT Train (2 epochs, ~25-40 min on T4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTConfig, SFTTrainer\n",
    "import inspect\n",
    "\n",
    "def safe_kwargs(cls, kw):\n",
    "    sig = inspect.signature(cls).parameters\n",
    "    return {k: v for k, v in kw.items() if k in sig}\n",
    "\n",
    "raw = dict(\n",
    "    output_dir='/content/opsguard-sft',\n",
    "    per_device_train_batch_size=2,\n",
    "    gradient_accumulation_steps=4,\n",
    "    num_train_epochs=2,\n",
    "    learning_rate=1e-4,\n",
    "    warmup_ratio=0.03,\n",
    "    logging_steps=5,\n",
    "    save_strategy='epoch',\n",
    "    save_total_limit=1,\n",
    "    bf16=True,\n",
    "    max_seq_length=2048,\n",
    "    dataset_text_field='text',\n",
    "    report_to='none',\n",
    "    push_to_hub=False,\n",
    ")\n",
    "cfg = SFTConfig(**safe_kwargs(SFTConfig, raw))\n",
    "trainer = SFTTrainer(model=model, train_dataset=ds, args=cfg, processing_class=tok)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Save + push to Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "api = HfApi(token=os.environ['HF_TOKEN'])\n",
    "api.create_repo('sai1906/opsguard-sft', repo_type='model', exist_ok=True)\n",
    "model.save_pretrained('/content/opsguard-sft-lora')\n",
    "tok.save_pretrained('/content/opsguard-sft-lora')\n",
    "api.upload_folder(folder_path='/content/opsguard-sft-lora', repo_id='sai1906/opsguard-sft', repo_type='model')\n",
    "print('Pushed: https://huggingface.co/sai1906/opsguard-sft')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Quick eval — trained vs base on E2_social_eng_buildup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, re, torch\n",
    "from server.opsguard_environment import OpsguardEnvironment\n",
    "from models import OpsguardAction, ActionType\n",
    "\n",
    "model.eval()\n",
    "SYSTEM_PROMPT = open('scripts/system_prompt.py').read().split('SYSTEM_PROMPT = ')[1].split('\"\"\"')[1]\n",
    "\n",
    "def format_obs(obs):\n",
    "    payload = {'scenario': obs.scenario_id, 'step': obs.step, 'feedback': obs.feedback,\n",
    "               'recent_actions': obs.recent_actions, 'memory_hits': obs.memory_hits}\n",
    "    if obs.current_issue:\n",
    "        ci = obs.current_issue\n",
    "        payload['current_issue'] = {'issue_id': ci.issue_id, 'number': ci.number,\n",
    "                                     'title': ci.title, 'body': ci.body[:600], 'is_pr': ci.is_pr,\n",
    "                                     'author_login': ci.author_login,\n",
    "                                     'pr_diff_preview': ci.pr_diff_preview[:500],\n",
    "                                     'available_labels': ci.available_labels[:20]}\n",
    "    return json.dumps(payload)\n",
    "\n",
    "def parse(text):\n",
    "    m = re.search(r'\\{[\\s\\S]*\\}', text)\n",
    "    if not m: return OpsguardAction(action_type=ActionType.WAIT)\n",
    "    try: data = json.loads(m.group(0))\n",
    "    except: return OpsguardAction(action_type=ActionType.WAIT)\n",
    "    try: return OpsguardAction(**data)\n",
    "    except: return OpsguardAction(action_type=ActionType.WAIT)\n",
    "\n",
    "@torch.inference_mode()\n",
    "def policy(obs):\n",
    "    prompt = SYSTEM_PROMPT + '\\n\\nOBSERVATION:\\n' + format_obs(obs) + '\\n\\nACTION:\\n'\n",
    "    inp = tok(prompt, return_tensors='pt', truncation=True, max_length=1800).to(model.device)\n",
    "    out = model.generate(**inp, max_new_tokens=200, do_sample=True, temperature=0.4,\n",
    "                         pad_token_id=tok.eos_token_id)\n",
    "    txt = tok.decode(out[0][inp['input_ids'].shape[1]:], skip_special_tokens=True)\n",
    "    return parse(txt)\n",
    "\n",
    "env = OpsguardEnvironment()\n",
    "obs = env.reset(scenario_id='E2_social_eng_buildup', seed=0)\n",
    "cum = 0; n_caught = 0\n",
    "for _ in range(40):\n",
    "    a = policy(obs)\n",
    "    obs = env.step(a)\n",
    "    if obs.reward: cum += obs.reward\n",
    "    if obs.done: break\n",
    "print(f'TRAINED REWARD on E2: {cum:+.2f}')\n",
    "print(f'final metadata: {obs.metadata}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}