organize_and_upload_trainers.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import json\n",
    "import datetime\n",
    "\n",
    "base_dir = \"full_sweep_v2\"\n",
    "new_dir = f\"{base_dir}_organized\"\n",
    "date = \"1230\"\n",
    "\n",
    "# We’ll loop through each folder under `base_dir`\n",
    "# that starts with \"resid_post_layer_*\".\n",
    "for folder in os.listdir(base_dir):\n",
    "    if folder.startswith(\"resid_post_layer_\"):\n",
    "        full_subdir_path = os.path.join(base_dir, folder)\n",
    "        \n",
    "        # Skip if it's not actually a directory\n",
    "        if not os.path.isdir(full_subdir_path):\n",
    "            continue\n",
    "        \n",
    "        # For each trainer directory (e.g. trainer_0, trainer_1...)\n",
    "        for trainer_dir in os.listdir(full_subdir_path):\n",
    "            trainer_path = os.path.join(full_subdir_path, trainer_dir)\n",
    "            \n",
    "            # Skip if it's not a folder\n",
    "            if not os.path.isdir(trainer_path):\n",
    "                continue\n",
    "            \n",
    "            # We expect a config.json inside each trainer directory\n",
    "            config_path = os.path.join(trainer_path, \"config.json\")\n",
    "            if not os.path.isfile(config_path):\n",
    "                continue  # If no config.json, skip\n",
    "            \n",
    "            # Read the config file\n",
    "            with open(config_path, \"r\") as f:\n",
    "                config = json.load(f)\n",
    "            \n",
    "            # Extract the trainer class\n",
    "            trainer_class = config[\"trainer\"][\"trainer_class\"]\n",
    "            model_name = config[\"trainer\"][\"lm_name\"]\n",
    "            ctx_len = config[\"buffer\"][\"ctx_len\"]\n",
    "\n",
    "            new_folder = f\"{trainer_class}_{model_name}_ctx{ctx_len}_{date}\".replace(\"/\", \"_\")\n",
    "            \n",
    "            # Construct the target path: <TrainerClass>/<resid_post_layer_X>/<trainer_#>\n",
    "            new_folder_path = os.path.join(new_dir, new_folder, folder)\n",
    "            new_trainer_path = os.path.join(new_folder_path, trainer_dir)\n",
    "            \n",
    "            # Make sure the new folder exists\n",
    "            os.makedirs(new_folder_path, exist_ok=True)\n",
    "            \n",
    "            # Copy (instead of move) the entire trainer directory\n",
    "            shutil.copytree(trainer_path, new_trainer_path, dirs_exist_ok=True)\n",
    "            \n",
    "            print(f\"Copied {trainer_path} -> {new_trainer_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import json\n",
    "import datetime\n",
    "\n",
    "base_dir = \"4k_gemma\"\n",
    "new_dir = f\"{base_dir}_organized\"\n",
    "date = \"0104\"\n",
    "\n",
    "# We’ll loop through each folder under `base_dir`\n",
    "# that starts with \"resid_post_layer_*\".\n",
    "for model_folder in os.listdir(base_dir):\n",
    "    for folder in os.listdir(os.path.join(base_dir, model_folder)):\n",
    "        if folder.startswith(\"resid_post_layer_\"):\n",
    "            full_subdir_path = os.path.join(base_dir, model_folder, folder)\n",
    "            \n",
    "            # Skip if it's not actually a directory\n",
    "            if not os.path.isdir(full_subdir_path):\n",
    "                continue\n",
    "\n",
    "            \n",
    "            # For each trainer directory (e.g. trainer_0, trainer_1...)\n",
    "            for trainer_dir in os.listdir(full_subdir_path):\n",
    "                trainer_path = os.path.join(full_subdir_path, trainer_dir)\n",
    "                \n",
    "                # Skip if it's not a folder\n",
    "                if not os.path.isdir(trainer_path):\n",
    "                    continue\n",
    "                \n",
    "                # We expect a config.json inside each trainer directory\n",
    "                config_path = os.path.join(trainer_path, \"config.json\")\n",
    "                if not os.path.isfile(config_path):\n",
    "                    continue  # If no config.json, skip\n",
    "                \n",
    "                # Read the config file\n",
    "                with open(config_path, \"r\") as f:\n",
    "                    config = json.load(f)\n",
    "                \n",
    "                # Extract the trainer class\n",
    "                trainer_class = config[\"trainer\"][\"trainer_class\"]\n",
    "                model_name = config[\"trainer\"][\"lm_name\"]\n",
    "                ctx_len = config[\"buffer\"][\"ctx_len\"]\n",
    "\n",
    "                new_folder = f\"{trainer_class}_{model_name}_ctx{ctx_len}_{date}\".replace(\"/\", \"_\")\n",
    "                \n",
    "                # Construct the target path: <TrainerClass>/<resid_post_layer_X>/<trainer_#>\n",
    "                new_folder_path = os.path.join(new_dir, new_folder, folder)\n",
    "                new_trainer_path = os.path.join(new_folder_path, trainer_dir)\n",
    "                \n",
    "                # Make sure the new folder exists\n",
    "                os.makedirs(new_folder_path, exist_ok=True)\n",
    "                \n",
    "                # Copy (instead of move) the entire trainer directory\n",
    "                shutil.copytree(trainer_path, new_trainer_path, dirs_exist_ok=True)\n",
    "                \n",
    "                print(f\"Copied {trainer_path} -> {new_trainer_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "api = HfApi()\n",
    "\n",
    "api.upload_folder(\n",
    "    folder_path=\"full_sweep_v2_organized\",\n",
    "    path_in_repo=\"\",\n",
    "    repo_id=\"adamkarvonen/sae_test\",\n",
    "    repo_type=\"model\",\n",
    "    # allow_patterns=\"*eval_results.json\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}