{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2fc22528",
   "metadata": {},
   "source": [
    "## Ay 119: LLMs - BERT and GPT \n",
    "\n",
    "Matthew Graham 2026\n",
    "\n",
    "These exercises will go through various aspects of the BERT and GPT architectures to give you a better feel for them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73496fa4",
   "metadata": {},
   "source": [
    "### Setup\n",
    "\n",
    "The notebook should use `transformers`, `sentence-transformers` and `openai`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51547158",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install transformers torch sentence-transformers openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc05b14e",
   "metadata": {},
   "source": [
    "## 1. Tokenization and Masked Language Modeling with BERT\n",
    "\n",
    "We will use a pretrained BERT tokenizer and masked language model on astronomy text.\n",
    "\n",
    "1. Tokenize astronomy terms such as:\n",
    "\n",
    "```text\n",
    "AGN, quasar, redshift, TDE, supernovae, H-alpha, Lambda-CDM, ZTF, JWST\n",
    "```\n",
    "\n",
    "2. Use BERT to fill in masked tokens in astronomy sentences such as:\n",
    "\n",
    "```text\n",
    "The galaxy has a redshift of [MASK].\n",
    "The broad H beta line is used to estimate the black hole [MASK].\n",
    "A tidal disruption event occurs when a star is disrupted by a massive black [MASK].\n",
    "Type Ia supernovae are useful as standard [MASK].\n",
    "The quasar was observed by [MASK] in the optical band.\n",
    "```\n",
    "\n",
    "Questions:\n",
    "1. Which astronomy terms are split into subwords?\n",
    "2. Does BERT predict sensible masked words?\n",
    "3. Where does BERT fail, and why?\n",
    "4. How might SciBERT, AstroBERT, or another domain-adapted model improve the result?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "757f190d",
   "metadata": {},
   "outputs": [],
   "source": [
    "terms = [\n",
    "    \"AGN\", \"quasar\", \"redshift\", \"TDE\", \"supernovae\",\n",
    "    \"H-alpha\", \"Lambda-CDM\", \"ZTF\", \"JWST\",\n",
    "]\n",
    "\n",
    "masked_sentences = [\n",
    "    \"The galaxy has a redshift of [MASK].\",\n",
    "    \"The broad H beta line is used to estimate the black hole [MASK].\",\n",
    "    \"A tidal disruption event occurs when a star is disrupted by a massive black [MASK].\",\n",
    "    \"Type Ia supernovae are useful as standard [MASK].\",\n",
    "    \"The quasar was observed by [MASK] in the optical band.\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8c8cda5",
   "metadata": {},
   "source": [
    "Have a look at using a pretrained tokenizer from Transformers with the bert-base-uncased model and the fill-mask pipeline for the MLM."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b228c180",
   "metadata": {},
   "source": [
    "## 2. BERT Embeddings for Astronomy Paper Similarity\n",
    "\n",
    "We are going to take 10–20 astronomy titles or short abstracts from different topics. Use BERT, SciBERT, or a sentence-transformer model to embed each text. Compute cosine similarities and find the nearest neighbors.\n",
    "\n",
    "Questions:\n",
    "1. Do papers on similar topics end up near each other?\n",
    "2. Which examples are grouped correctly?\n",
    "3. Which examples are surprising or wrong?\n",
    "4. Why is BERT useful for retrieval/classification even though it is not generating text?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83f2ab9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = [\n",
    "    \"A luminous optical flare from a tidal disruption event in the nucleus of an active galaxy\",\n",
    "    \"Mid-infrared echoes from dust reprocessing in tidal disruption events\",\n",
    "    \"Broad emission line variability and black hole mass estimates in quasars\",\n",
    "    \"Changing-look quasars and accretion state transitions in active galactic nuclei\",\n",
    "    \"Type Ia supernova cosmology with improved standard candle calibration\",\n",
    "    \"Core-collapse supernova light curves from massive stellar explosions\",\n",
    "    \"Transmission spectroscopy of a warm Neptune exoplanet atmosphere with JWST\",\n",
    "    \"Transit timing variations reveal an additional planet in a compact exoplanet system\",\n",
    "    \"Galaxy stellar mass functions and quenching across cosmic time\",\n",
    "    \"The evolution of star-forming galaxies in deep optical surveys\",\n",
    "    \"Constraints on dark energy from baryon acoustic oscillations and supernova distances\",\n",
    "    \"Weak lensing measurements of matter clustering in a large cosmological survey\",\n",
    "]\n",
    "\n",
    "labels = [\n",
    "    \"TDE\", \"TDE\", \"AGN\", \"AGN\", \"Supernova\", \"Supernova\",\n",
    "    \"Exoplanet\", \"Exoplanet\", \"Galaxy\", \"Galaxy\", \"Cosmology\", \"Cosmology\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14767751",
   "metadata": {},
   "source": [
    "## 3. GPT Prompting for Astronomy Abstracts and Structured Claim Audits\n",
    "\n",
    "We are going to use GPT as a conditional text generator, but require it to produce output in a structure that can be checked programmatically. The idea is to treat GPT output as data that can be audited.\n",
    "\n",
    "1. Build prompts from a small event metadata dictionary.\n",
    "2. Call GPT.\n",
    "3. Ask for both prose and JSON-like claim summaries.\n",
    "4. Automatically flag numbers, object names, and claims that require checking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef0285a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "event = {\n",
    "    \"object_name\": \"AT2025demo\",\n",
    "    \"host_type\": \"AGN\",\n",
    "    \"redshift\": 0.145,\n",
    "    \"peak_brightening_mag\": 1.8,\n",
    "    \"bands\": [\"ZTF g\", \"ZTF r\", \"WISE W1\", \"WISE W2\"],\n",
    "    \"spectroscopy\": \"broad Balmer emission and variable blue continuum\",\n",
    "    \"mid_ir_lag_days\": 210,\n",
    "    \"interpretations_to_discuss\": [\n",
    "        \"tidal disruption event in an AGN disk\",\n",
    "        \"changing-look accretion-state transition\",\n",
    "        \"extreme stochastic AGN variability\",\n",
    "    ],\n",
    "}\n",
    "\n",
    "def format_event(ev):\n",
    "    return \"\\n\".join(f\"- {k}: {v}\" for k, v in ev.items())\n",
    "\n",
    "print(format_event(event))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ed7f3ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "def call_gpt(prompt, *, system=None, model=GPT_MODEL, temperature=0.2, max_output_tokens=700, mode=\"minimal\"):\n",
    "    \"\"\"Call GPT via the OpenAI SDK \n",
    "\n",
    "    Notes:\n",
    "    - Never hard-code API keys in notebooks.\n",
    "    - Keep temperature low for scientific summarization/evaluation tasks.\n",
    "    - Treat the result as a draft or hypothesis, not as a source of truth.\n",
    "    \"\"\"\n",
    "    system = system or \"You are a careful astronomy research assistant. Distinguish observations from interpretations.\"\n",
    "\n",
    "    try:\n",
    "        client = OpenAI()\n",
    "        response = client.responses.create(\n",
    "            model=model,\n",
    "            input=[\n",
    "                {\"role\": \"system\", \"content\": system},\n",
    "                {\"role\": \"user\", \"content\": prompt},\n",
    "            ],\n",
    "            temperature=temperature,\n",
    "            max_output_tokens=max_output_tokens,\n",
    "        )\n",
    "        return response.output_text\n",
    "    except Exception as exc:\n",
    "        print(f\"Live GPT call failed; Error: {exc}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84e40971",
   "metadata": {},
   "source": [
    "## 4. GPT as an Astronomy Coding Assistant with Automated Tests\n",
    "\n",
    "We will use GPT to write a small astronomy utility function, then evaluate the generated code with tests. In a \n",
    "real workflow, GPT might propose code, but a human scientist plus an automated evaluator would decide whether \n",
    "it is acceptable.\n",
    "\n",
    "You will ask GPT to write code for AB magnitude/flux conversions and then test the result against known values.\n",
    "1. Extract and safety check the generated code\n",
    "2. Run unit tests against the GPT-generated code\n",
    "3. Prompt code revision based on failing tests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aae4873f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}