{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro",
   "metadata": {},
   "source": [
    "# Course 3 — Post-Training and Agentic Learning\n",
    "Executable teaching miniatures for Kimi K3 §4. **No policy training, reward-model evaluation, MXFP kernel, or paper-scale environment is reproduced.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bootstrap",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T01:18:00.017620Z",
     "iopub.status.busy": "2026-08-11T01:18:00.017121Z",
     "iopub.status.idle": "2026-08-11T01:18:01.021743Z",
     "shell.execute_reply": "2026-08-11T01:18:01.020078Z"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import subprocess, sys, tempfile, torch\n",
    "repo = next((p for p in (Path.cwd(), *Path.cwd().parents) if (p / 'src').is_dir()), None)\n",
    "if repo is None:\n",
    "    repo = (Path('/content') if Path('/content').is_dir() else Path(tempfile.gettempdir())) / 'build-Kimi-K3-architecture'\n",
    "    if not (repo / 'src').is_dir():\n",
    "        subprocess.run(['git', 'clone', '--depth', '1', 'https://github.com/mailtotanvir/build-Kimi-K3-architecture.git', str(repo)], check=True)\n",
    "sys.path.insert(0, str(repo))\n",
    "from src.posttraining.multi_effort_rl import effort_reward, pause_at_fraction\n",
    "from src.posttraining.multi_teacher_distill import MultiTeacherDistillationLoss\n",
    "from src.posttraining.deployment_aware import speculative_acceptance_probability, lk_loss\n",
    "from src.posttraining.kg_task_synthesis import KnowledgeGraphTaskSynthesizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "partial-rollout",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T01:18:01.023583Z",
     "iopub.status.busy": "2026-08-11T01:18:01.023370Z",
     "iopub.status.idle": "2026-08-11T01:18:01.027345Z",
     "shell.execute_reply": "2026-08-11T01:18:01.026508Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "completed at λ=.5: ['search-1', 'code-1']\n",
      "paused for next iteration: ['search-2', 'assistant-1']\n"
     ]
    }
   ],
   "source": [
    "finish = {'search-1':2,'search-2':11,'code-1':4,'assistant-1':15}\n",
    "batch = pause_at_fraction({k:0 for k in finish}, finish, .5)\n",
    "print('completed at λ=.5:', batch.completed)\n",
    "print('paused for next iteration:', batch.paused)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "effort",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T01:18:01.029224Z",
     "iopub.status.busy": "2026-08-11T01:18:01.029054Z",
     "iopub.status.idle": "2026-08-11T01:18:01.032980Z",
     "shell.execute_reply": "2026-08-11T01:18:01.031845Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "used= 70, reward=+0.8\n",
      "used=100, reward=+0.8\n",
      "used=130, reward=-1.0\n",
      "The 100-token fixture is illustrative; the paper does not publish fixed effort budgets.\n"
     ]
    }
   ],
   "source": [
    "for used in (70,100,130):\n",
    "    print(f'used={used:3d}, reward={effort_reward(.8,used,100,1.0):+.1f}')\n",
    "print('The 100-token fixture is illustrative; the paper does not publish fixed effort budgets.')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "mopd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T01:18:01.034654Z",
     "iopub.status.busy": "2026-08-11T01:18:01.034503Z",
     "iopub.status.idle": "2026-08-11T01:18:01.039469Z",
     "shell.execute_reply": "2026-08-11T01:18:01.038175Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eq. 15 clipped token rewards: [0.6931, -0.6931, 0.6931]\n"
     ]
    }
   ],
   "source": [
    "student = torch.tensor([.20,.50,.10])\n",
    "teacher = torch.tensor([.40,.25,.20])\n",
    "rewards = MultiTeacherDistillationLoss(2.0).token_rewards(student,teacher)\n",
    "print('Eq. 15 clipped token rewards:', [round(x,4) for x in rewards.tolist()])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "lk",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T01:18:01.041718Z",
     "iopub.status.busy": "2026-08-11T01:18:01.041554Z",
     "iopub.status.idle": "2026-08-11T01:18:01.046367Z",
     "shell.execute_reply": "2026-08-11T01:18:01.045397Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acceptance overlap: 0.8\n",
      "LK loss: 0.2231\n"
     ]
    }
   ],
   "source": [
    "p=torch.tensor([[.6,.3,.1]])\n",
    "q=torch.tensor([[.5,.2,.3]])\n",
    "print('acceptance overlap:', round(speculative_acceptance_probability(p,q).item(),4))\n",
    "print('LK loss:', round(lk_loss(p,q).item(),4))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "kg",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T01:18:01.047967Z",
     "iopub.status.busy": "2026-08-11T01:18:01.047823Z",
     "iopub.status.idle": "2026-08-11T01:18:01.051770Z",
     "shell.execute_reply": "2026-08-11T01:18:01.050697Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "valid DAG: True\n",
      "synthetic task: Explain how AI connects to Attention -> Attention connects to KDA\n"
     ]
    }
   ],
   "source": [
    "kg=KnowledgeGraphTaskSynthesizer(['AI','Attention','KDA'], [('AI','Attention'),('Attention','KDA')])\n",
    "print('valid DAG:',kg.validate_dag())\n",
    "print('synthetic task:',kg.synthesize_task()['prompt'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "boundary",
   "metadata": {},
   "source": [
    "## Evidence boundary\n",
    "Executed: queue cutoff/resumption, effort penalty, Equation 15 log-ratio reward, Equation 16 overlap loss, and DAG validation.\n",
    "\n",
    "Paper reported: nine expert policies, partial-rollout stability, Agentic GRM, full MOPD, MXFP QAT, EAGLE-3 speedup, and all production environments."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}