try_langchain_tune_ea.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.append(\"ares_transverse_tuning\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "\n",
    "import gymnasium as gym\n",
    "import langchain\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scienceplots\n",
    "from dotenv import load_dotenv\n",
    "from gymnasium.wrappers import TimeLimit\n",
    "from icecream import ic\n",
    "from langchain.callbacks import FileCallbackHandler, wandb_tracing_enabled\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.chat_models import ChatOllama\n",
    "from langchain.output_parsers import ResponseSchema, StructuredOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "from loguru import logger\n",
    "from openai import RateLimitError\n",
    "from src.environments.ea import TransverseTuning\n",
    "from src.eval import Episode\n",
    "from src.trial import load_trials\n",
    "from src.wrappers import RecordEpisode, TQDMWrapper\n",
    "\n",
    "from pacuna import PACuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "\n",
    "plt.style.use([\"science\", \"nature\", \"no-latex\"])\n",
    "\n",
    "langchain.debug = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# chat_model = ChatOpenAI(model=\"gpt-3.5-turbo-0125\")\n",
    "# chat_model = ChatOpenAI(model=\"gpt-4\")\n",
    "# chat_model = ChatOpenAI(model=\"gpt-4-32k\")\n",
    "# chat_model = ChatOpenAI(model=\"gpt-4-0125-preview\")\n",
    "# chat_model = ChatOllama(model=\"mistral:v0.2\", base_url=\"http://max-wng053:11434\")\n",
    "# chat_model = ChatOllama(model=\"zephyr\", base_url=\"http://max-wng052:11434\")\n",
    "# chat_model = ChatOllama(model=\"mixtral:8x7b\", base_url=\"http://max-cssbg012:11434\")\n",
    "# chat_model = ChatOllama(model=\"gemma:2b\", base_url=\"http://max-wng053:11434\")\n",
    "# chat_model = ChatOllama(model=\"gemma:7b\", base_url=\"http://max-wng053:11434\")\n",
    "# chat_model = ChatOllama(\n",
    "#     model=\"starling-lm:7b-beta\", base_url=\"http://max-cssbg012:11434\"\n",
    "# )\n",
    "chat_model = ChatOllama(model=\"vicuna:7b-16k\", base_url=\"http://max-cssbg012:11434\")\n",
    "# chat_model = ChatOllama(model=\"vicuna:33b\", base_url=\"http://max-wng054:11434\")\n",
    "# chat_model = ChatOllama(model=\"llava:34b\", base_url=\"http://max-wng054:11434\")\n",
    "# chat_model = ChatOllama(model=\"orca2:7b\", base_url=\"http://max-cssbg012:11434\")\n",
    "# chat_model = ChatOllama(model=\"orca2:13b\", base_url=\"http://max-cssbg012:11434\")\n",
    "# chat_model = ChatOllama(model=\"llama2:7b\", base_url=\"http://max-wng053:11434\")\n",
    "# chat_model = ChatOllama(model=\"llama2:13b\", base_url=\"http://max-wng053:11434\")\n",
    "# chat_model = ChatOllama(model=\"llama2:70b\", base_url=\"http://max-wng053:11434\")\n",
    "# chat_model = ChatOllama(model=\"falcon:180b-chat\", base_url=\"http://max-wng054:11434\")\n",
    "# chat_model = ChatOllama(model=\"neural-chat:7b\")\n",
    "# chat_model = ChatOllama(model=\"mistral-openorca:7b\", base_url=\"http://max-wng052:11434\")\n",
    "# chat_model = ChatOllama(model=\"phi:chat\")\n",
    "# chat_model = ChatOllama(model=\"megadolphin:120b\", base_url=\"http://max-wng053:11434\")\n",
    "# chat_model = ChatOllama(model=\"yi:34b-chat\", base_url=\"http://max-wng054:11434\") x\n",
    "# chat_model = PACuna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# chat_model.invoke(\"Hello\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRIAL_INDEX = 38\n",
    "trials = load_trials(Path(\"ares_transverse_tuning/data/trials.yaml\"))\n",
    "trial = trials[TRIAL_INDEX]\n",
    "trial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = (\n",
    "    chat_model.model_name if hasattr(chat_model, \"model_name\") else chat_model.model\n",
    ").replace(\":\", \"-\")\n",
    "now = datetime.now()\n",
    "\n",
    "log_dir = (\n",
    "    Path(\"data\")\n",
    "    / \"paper\"\n",
    "    / \"explained\"\n",
    "    / model_name\n",
    "    / f\"trial-{TRIAL_INDEX}_{now.strftime('%Y-%m-%d_%H-%M-%S')}\"\n",
    ")\n",
    "langchain_log_file = log_dir / \"langchain.log\"\n",
    "logger.add(langchain_log_file, colorize=True, enqueue=True)\n",
    "handler = FileCallbackHandler(langchain_log_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = TransverseTuning(\n",
    "    action_mode=\"direct\",\n",
    "    magnet_init_mode=trial.initial_magnets,\n",
    "    target_beam_mode=trial.target_beam,\n",
    "    backend=\"cheetah\",\n",
    "    backend_args={\n",
    "        \"incoming_mode\": trial.incoming_beam,\n",
    "        \"misalignment_mode\": trial.misalignments,\n",
    "        \"generate_screen_images\": False,\n",
    "    },\n",
    ")\n",
    "# env = TimeLimit(env, max_episode_steps=18)\n",
    "# env = RecordEpisode(env, Path(\"data/recorded_episodes\"))\n",
    "# env = TQDMWrapper(env)\n",
    "env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample_template = \"\"\"Magnet settings:\n",
    "# ```json\n",
    "# {{\n",
    "# \t\"Q1\": {q1:.2f},  // m^-2\n",
    "# \t\"Q2\": {q2:.2f},  // m^-2\n",
    "# \t\"CV\": {cv:.2f},  // mrad\n",
    "# \t\"Q3\": {q3:.2f},  // m^-2\n",
    "# \t\"CH\": {ch:.2f}  // mrad\n",
    "# }}\n",
    "# ```\n",
    "# Beam parameters:\n",
    "# ```json\n",
    "# {{\n",
    "# \t\"horizontal position\": {mu_x:.2f},  // mm\n",
    "# \t\"horizontal size\": {sigma_x:.2f},  // mm\n",
    "# \t\"vertical position\": {mu_y:.2f},  // mm\n",
    "# \t\"vertical size\": {sigma_y:.2f}  // mm\n",
    "# }}\n",
    "# ```\n",
    "# \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_template = \"\"\"Magnet settings:\n",
    "```json\n",
    "{{\n",
    "\t\"Q1\": {q1:.2f},\n",
    "\t\"Q2\": {q2:.2f},\n",
    "\t\"CV\": {cv:.2f},\n",
    "\t\"Q3\": {q3:.2f},\n",
    "\t\"CH\": {ch:.2f}\n",
    "}}\n",
    "```\n",
    "Beam parameters:\n",
    "```json\n",
    "{{\n",
    "\t\"mu_x\": {mu_x:.2f},\n",
    "\t\"sigma_x\": {sigma_x:.2f},\n",
    "\t\"mu_y\": {mu_y:.2f},\n",
    "\t\"sigma_y\": {sigma_y:.2f}\n",
    "}}\n",
    "```\n",
    "\"\"\"\n",
    "target_template = \"\"\"Target beam parameters:\n",
    "```json\n",
    "{{\n",
    "\t\"mu_x\": {mu_x:.2f},\n",
    "\t\"sigma_x\": {sigma_x:.2f},\n",
    "\t\"mu_y\": {mu_y:.2f},\n",
    "\t\"sigma_y\": {sigma_y:.2f}\n",
    "}}\n",
    "```\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "message_template = \"\"\"Now you will help me optimise the horizontal and vertical position and size of an electron beam on a diagnostic screen in a particle accelerator.\n",
    "\n",
    "You are able to control five magnets in the beam line. The magnets are called Q1, Q2, CV, Q3 and CH.\n",
    "\n",
    "Q1, Q2 and Q3 are quadrupole magnets. When their k1 strength is increased, the beam becomes more focused in the horizontal plane and more defocused in the vertical plane. When their k1 strength is decreased, the beam becomes more focused in the vertical plane and more defocused in the horizontal plane. When their k1 strength is zero, the beam is not focused in either plane. Quadrupole magnets might also steer the beam in the horizontal or vertical plane depending on their k0 strength, when the beam does not travel through the centre of the magnet. The range of the k1 strength is -30.0 to 30.0 m^-2.\n",
    "\n",
    "CV is vertical steering magnet. When its deflection angle is increased, the beam is steered upwards. When its deflection angle is decreased, the beam is steered downwards. The range of the deflection angle is -6.0 to 6.0 mrad.\n",
    "\n",
    "CH is horizontal steering magnet. When its deflection angle is increased, the beam is steered to the right. When its deflection angle is decreased, the beam is steered to the left. The range of the deflection angle is -6.0 to 6.0 mrad.\n",
    "\n",
    "You are optimising four beam parameters: mu_x, sigma_x, mu_y, sigma_y. The beam parameters are measured in millimetres (mm). The target beam parameters are:\n",
    "\n",
    "{target_beam_parameters}\n",
    "\n",
    "Below are previously measured pairs of magnet settings and the corresponding observed beam parameters.\n",
    "\n",
    "{prior_samples}\n",
    "\n",
    "Give me new magnet settings that are different from all pairs above. The magnet settings you should propose should lead to beam parameters closer the target or, if you do not have enough information yet, maximise information gain for finding new beam parameters. Do not set any magnet setting to zero. Smooth changes relative to the last magnet settings are preferred.\n",
    "\n",
    "{format_instructions}\n",
    "\n",
    "Do not add comments to the output JSON.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_q1 = ResponseSchema(\n",
    "    name=\"Q1\",\n",
    "    type=\"float\",\n",
    "    description=\"k1 strength of the first quadrupole magnet\",\n",
    ")\n",
    "response_q2 = ResponseSchema(\n",
    "    name=\"Q2\",\n",
    "    type=\"float\",\n",
    "    description=\"k1 strength of the second quadrupole magnet\",\n",
    ")\n",
    "response_cv = ResponseSchema(\n",
    "    name=\"CV\",\n",
    "    type=\"float\",\n",
    "    description=\"Deflection angle of the vertical steering magnet\",\n",
    ")\n",
    "response_q3 = ResponseSchema(\n",
    "    name=\"Q3\",\n",
    "    type=\"float\",\n",
    "    description=\"k1 strength of the third quadrupole magnet\",\n",
    ")\n",
    "response_ch = ResponseSchema(\n",
    "    name=\"CH\",\n",
    "    type=\"float\",\n",
    "    description=\"Deflection angle of the horizontal steering magnet\",\n",
    ")\n",
    "\n",
    "output_parser = StructuredOutputParser.from_response_schemas(\n",
    "    [response_q1, response_q2, response_cv, response_q3, response_ch]\n",
    ")\n",
    "\n",
    "print(output_parser.get_format_instructions())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = PromptTemplate.from_template(message_template)\n",
    "prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(\n",
    "    prompt=prompt, llm=chat_model, output_parser=output_parser, callbacks=[handler]\n",
    ")\n",
    "decapitated_chain = LLMChain(prompt=prompt, llm=chat_model)\n",
    "chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # with wandb_tracing_enabled():\n",
    "# response = decapitated_chain.invoke(\n",
    "#     {\n",
    "#         \"prior_samples\": samples_str,\n",
    "#         \"format_instructions\": output_parser.get_format_instructions(),\n",
    "#     }\n",
    "# )\n",
    "# print(response[\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with wandb_tracing_enabled():\n",
    "#     response = chain.invoke(\n",
    "#         {\n",
    "#             \"prior_samples\": samples_str,\n",
    "#             \"format_instructions\": output_parser.get_format_instructions(),\n",
    "#         }\n",
    "#     )\n",
    "# print(response[\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LLMTransverseTuningAgent:\n",
    "    \"\"\"\n",
    "    An agent for doing transverse beam parameter tuning in the ARES experimental\n",
    "    area using an LLM via LangChain.\n",
    "    \"\"\"\n",
    "\n",
    "    #     message_template = \"\"\"Now you will help me optimise the horizontal and vertical position and size of an electron beam on a diagnostic screen in a particle accelerator.\n",
    "    #\n",
    "    # The target beam parameters I want you to find are:\n",
    "    #  - horizontal position: 0.0 μm\n",
    "    #  - horizontal size: 0.0 μm\n",
    "    #  - vertical position: 0.0 μm\n",
    "    #  - vertical size: 0.0 μm\n",
    "    #\n",
    "    # You are able to control five magnets in the beam line. The magnets are called:\n",
    "    #  - Q1\n",
    "    #  - Q2\n",
    "    #  - CV\n",
    "    #  - Q3\n",
    "    #  - CH\n",
    "    #\n",
    "    # Q1, Q2, Q3 are quadrupole magnets. When their k1 strenth is increased, the beam becomes more focused in the horizontal plane and more defocused in the vertical plane. When their k1 strength is decreased, the beam becomes more focused in the vertical plane and more defocused in the horizontal plane. When their k1 strength is zero, the beam is not focused in either plane. Quadrupole magnets might also steer the beam in the horizontal or vertical plane depending on their k0 strength, when the beam does not travel through the centre of the magnet. The range of the k1 strength is -30.0 to 30.0 m^-2.\n",
    "    #\n",
    "    # CV is vertical steering magnet. When its deflection angle is increased, the beam is steered upwards. When its deflection angle is decreased, the beam is steered downwards. The range of the deflection angle is -6.0 to 6.0 mrad.\n",
    "    #\n",
    "    # CH is horizontal steering magnet. When its deflection angle is increased, the beam is steered to the right. When its deflection angle is decreased, the beam is steered to the left. The range of the deflection angle is -6.0 to 6.0 mrad.\n",
    "    #\n",
    "    # I have some pairs of magnet settings and the corresponding beam parameters.\n",
    "    #\n",
    "    # {prior_samples}\n",
    "    #\n",
    "    # Give me new magnet settings that are different from all pairs above, and will result in transnverse beam parameters closer to the target beam parameters than any of the above. Beam parameters less than 40 μm from their target are considered optimal. If you do not know which magnet settings would improve the beam parameters, choose magnet settings that maximise information gain. Smooth changes to the magnet settings are preferred. Do not write code.\n",
    "    #\n",
    "    # {format_instructions}\n",
    "    # \"\"\"\n",
    "\n",
    "    #     sample_template = \"\"\"Magnet settings:\n",
    "    #  - Q1: {q1} m^-2\n",
    "    #  - Q2: {q2} m^-2\n",
    "    #  - CV: {cv} mrad\n",
    "    #  - Q3: {q3} m^-2\n",
    "    #  - CH: {ch} mrad\n",
    "    # Beam parameters:\n",
    "    #  - horizontal position: {mu_x} μm\n",
    "    #  - horizontal size: {sigma_x} μm\n",
    "    #  - vertical position: {mu_y} μm\n",
    "    #  - vertical size: {sigma_y} μm\n",
    "    # \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, env: gym.Env, warmup_steps: int = 0, verbose: bool = False\n",
    "    ) -> None:\n",
    "        self.env = env\n",
    "        self.warmup_steps = warmup_steps\n",
    "        self.verbose = verbose\n",
    "\n",
    "        global chain\n",
    "        self._chain = chain\n",
    "\n",
    "        self._observations = []\n",
    "\n",
    "    def predict(self, observation: dict) -> np.ndarray:\n",
    "        \"\"\"\n",
    "        Takes an observation from the environment and returns an action.\n",
    "        \"\"\"\n",
    "        self._observations.append(observation)\n",
    "\n",
    "        # -1 because we already have the first observation from the reset\n",
    "        if len(self._observations) < self.warmup_steps - 1:\n",
    "            return self.env.action_space.sample()\n",
    "\n",
    "        global sample_template\n",
    "        global target_template\n",
    "        global output_parser\n",
    "\n",
    "        samples_str = \"\\n\".join(\n",
    "            sample_template.format(\n",
    "                q1=obs[\"magnets\"][0],\n",
    "                q2=obs[\"magnets\"][1],\n",
    "                cv=obs[\"magnets\"][2] * 1e3,\n",
    "                q3=obs[\"magnets\"][3],\n",
    "                ch=obs[\"magnets\"][4] * 1e3,\n",
    "                mu_x=obs[\"beam\"][0] * 1e6,\n",
    "                sigma_x=obs[\"beam\"][1] * 1e6,\n",
    "                mu_y=obs[\"beam\"][2] * 1e6,\n",
    "                sigma_y=obs[\"beam\"][3] * 1e6,\n",
    "            )\n",
    "            for obs in self._observations\n",
    "        )\n",
    "        target_str = target_template.format(\n",
    "            mu_x=observation[\"target\"][0] * 1e3,\n",
    "            sigma_x=observation[\"target\"][1] * 1e3,\n",
    "            mu_y=observation[\"target\"][2] * 1e3,\n",
    "            sigma_y=observation[\"target\"][3] * 1e3,\n",
    "        )\n",
    "\n",
    "        try:\n",
    "            response = self._chain.invoke(\n",
    "                {\n",
    "                    \"target_beam_parameters\": target_str,\n",
    "                    \"prior_samples\": samples_str,\n",
    "                    \"format_instructions\": output_parser.get_format_instructions(),\n",
    "                }\n",
    "            )\n",
    "        except Exception as e:\n",
    "            time.sleep(5.0)\n",
    "            response = self._chain.invoke(\n",
    "                {\n",
    "                    \"target_beam_parameters\": target_str,\n",
    "                    \"prior_samples\": samples_str,\n",
    "                    \"format_instructions\": output_parser.get_format_instructions(),\n",
    "                }\n",
    "            )\n",
    "\n",
    "        if self.verbose:\n",
    "            print(response[\"explanation\"])\n",
    "\n",
    "        action = np.array(\n",
    "            [\n",
    "                response[\"text\"][\"Q1\"],\n",
    "                response[\"text\"][\"Q2\"],\n",
    "                response[\"text\"][\"CV\"] / 1e3,\n",
    "                response[\"text\"][\"Q3\"],\n",
    "                response[\"text\"][\"CH\"] / 1e3,\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        return action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = TransverseTuning(\n",
    "    action_mode=\"direct\",\n",
    "    magnet_init_mode=trial.initial_magnets,\n",
    "    target_beam_mode=trial.target_beam,\n",
    "    backend=\"cheetah\",\n",
    "    backend_args={\n",
    "        \"incoming_mode\": trial.incoming_beam,\n",
    "        \"misalignment_mode\": trial.misalignments,\n",
    "        \"generate_screen_images\": False,\n",
    "    },\n",
    ")\n",
    "env = TimeLimit(env, max_episode_steps=50)\n",
    "env = RecordEpisode(env, log_dir / \"recorded_episodes\")\n",
    "env = TQDMWrapper(env)\n",
    "env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = LLMTransverseTuningAgent(env=env, warmup_steps=0, verbose=False)\n",
    "agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with wandb_tracing_enabled():\n",
    "    observation, info = env.reset()\n",
    "    done = False\n",
    "    while not done:\n",
    "        action = agent.predict(observation)\n",
    "        observation, _, terminated, truncated, _ = env.step(action)\n",
    "        done = terminated or truncated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "episode = Episode.load(log_dir / \"recorded_episodes\" / \"recorded_episode_1.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = episode.plot_summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm-accelerator-tuning",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}