more flexible TrainablePolicyModel

Author: corbt

dev/yes-no-maybe.ipynb

@@ -41,12 +41,12 @@
     "load_dotenv()\n",
     "\n",
     "\n",
-    "api = art.LocalAPI()\n",
-    "model = await api.get_or_create_model(\n",
+    "model = art.TrainablePolicyModel(\n",
     "    name=\"001\",\n",
     "    project=\"yes-no-maybe\",\n",
     "    base_model=\"Qwen/Qwen2.5-7B-Instruct\",\n",
     ")\n",
+    "await model.register_for_training(art.LocalAPI())\n",
     "\n",
     "\n",
     "async def rollout(client: openai.AsyncOpenAI, prompt: str) -> art.Trajectory:\n",
@@ -90,7 +90,7 @@
     "    ]\n",
     "]\n",
     "\n",
-    "openai_client = await model.openai_client()\n",
+    "openai_client = model.openai_client()\n",
     "for _ in range(await model.get_step(), 1_000):\n",
     "    train_groups = await art.gather_trajectory_groups(\n",
     "        (\n",

examples/2048/2048-single-turn.ipynb

@@ -47,10 +47,10 @@
     "random.seed(42)\n",
     "\n",
     "\n",
-    "api = art.LocalAPI()\n",
-    "model = await api.get_or_create_model(\n",
+    "model = art.TrainablePolicyModel(\n",
     "    name=\"001\", project=\"2048-single-turn\", base_model=\"Qwen/Qwen2.5-7B-Instruct\"\n",
-    ")"
+    ")\n",
+    "await model.register_for_training(art.LocalAPI())"
    ]
   },
   {
@@ -78,7 +78,6 @@
     "\n",
     "@art.retry(exceptions=(openai.LengthFinishReasonError, requests.ReadTimeout))\n",
     "async def rollout(client: openai.AsyncOpenAI, iteration: int) -> art.Trajectory:\n",
-    "\n",
     "    game = generate_game()\n",
     "\n",
     "    reward = 0\n",
@@ -87,7 +86,6 @@
     "    trajectories: list[art.Trajectory] = []\n",
     "\n",
     "    while True:\n",
-    "\n",
     "        trajectory = art.Trajectory(\n",
     "            messages_and_choices=[\n",
     "                {\n",
@@ -193,7 +191,7 @@
     "    return trajectories\n",
     "\n",
     "\n",
-    "openai_client = await model.openai_client()\n",
+    "openai_client = model.openai_client()\n",
     "\n",
     "for i in range(await model.get_step(), 500):\n",
     "    train_groups = await art.gather_trajectory_groups(\n",