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.TrainableModel(\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",