{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-000"
      },
      "source": "# CS145 NanoQuery\n\n**Starter release: October 19, 2026.**\n\n[NanoQuery: requirements and grading](https://cs145.web.app/projects/project2-science.html)\n\nConnect tables to documents and answer analytical questions. At higher tiers, build a tested SQL library and verified answers, then test whether giving an LLM that library improves generated SQL. No model training or fine-tuning is required.\n\n**Using the project page and notebook:** We’ve aimed to keep the project page and notebook consistent. The project page gives the full requirements and grading; the notebook provides compact guidance, function contracts, and starter code.\n\nThe supplied helpers support setup and retrieval; you write your dataset-specific queries, mappings, and tests through T87. For T93–T100, extend your work with verified reference answers and the LLM comparison. You may need to adapt the helpers to your implementation.\n\nYou’re welcome to adapt the starter code within the project requirements. If something seems inconsistent or buggy, check the project page and ask privately on Ed; include the error and a small reproducible example for bugs. Note important assumptions in your notebook so we can take them into account when reviewing your work.\n\nYour proposed tier is a planning target, not a score cap. We award partial credit within each rubric category; 80 core points plus 4 library points earns 84.\n\n**Bonus:** Staff select the top five Science projects for overall quality and compelling work; each receives 0–5 bonus points.\n\nRun Setup and Configuration first; the supplied helpers need no edits. Then jump to [Phase 0](#scrollTo=phase-0) or [Phase 1](#scrollTo=phase-1).\n",
      "id": "science-v2-000"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "heading_collapsed": true,
        "id": "science-v2-001"
      },
      "source": "## Setup\n\nFirst, finalize the Project 1 tables you will use for Project 2 (at least three tables totaling 500 MB) and select 100–200 related Wikipedia articles or supply your own relevant domain documents. You may choose a different public dataset.\n\nRun setup once, fill in Configuration, then run Connect to BigQuery and verify all three access checks. Keep source tables in BigQuery; no table download is needed. In Phase 0, prepare the articles, map them to table records, and complete your proposal.\n\nSpark needs Java 17 or 21. The setup cell checks Java and installs Java 17 in Colab if needed.\n",
      "id": "science-v2-001"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-002"
      },
      "source": "# @title Install dependencies\n%pip -q install pyspark==4.0.1 sentence-transformers google-cloud-bigquery pandas pyarrow db-dtypes sqlglot==30.22.0 requests\nimport os, json, time, re, random, platform\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nimport importlib.metadata as metadata\nprint(\"Python:\", platform.python_version())\nprint({p: metadata.version(p) for p in [\"pyspark\", \"sentence-transformers\", \"transformers\", \"torch\", \"google-cloud-bigquery\", \"sqlglot\"]})\nWORK = Path(\"p2_work\")\nWORK.mkdir(exist_ok=True)\n\n# Spark 4 requires Java 17 or 21. Configure Java before starting the JVM.\nimport shutil, subprocess, sys\n\ndef ensure_spark_java():\n    java_home = os.environ.get(\"JAVA_HOME\")\n    java = str(Path(java_home) / \"bin\" / \"java\") if java_home else shutil.which(\"java\")\n    major = None\n    if java:\n        try:\n            result = subprocess.run([java, \"-version\"], capture_output=True, text=True, check=True)\n            match = re.search(r'version \"(\\d+)', result.stderr + result.stdout)\n            major = int(match.group(1)) if match else None\n        except (OSError, subprocess.CalledProcessError):\n            pass\n    if major not in (17, 21):\n        try:\n            import google.colab\n        except ImportError:\n            raise RuntimeError(\"Spark needs Java 17 or 21. Install one and set JAVA_HOME, then rerun Setup.\") from None\n        subprocess.run([\"apt-get\", \"update\", \"-qq\"], check=True)\n        subprocess.run([\"apt-get\", \"install\", \"-y\", \"-qq\", \"openjdk-17-jdk-headless\"], check=True)\n        homes = sorted(Path(\"/usr/lib/jvm\").glob(\"java-17-openjdk-*\"))\n        if not homes:\n            raise RuntimeError(\"Java 17 installation was not found under /usr/lib/jvm.\")\n        java = str(homes[0] / \"bin\" / \"java\")\n    os.environ[\"JAVA_HOME\"] = str(Path(java).resolve().parent.parent)\n    subprocess.run([str(Path(os.environ[\"JAVA_HOME\"]) / \"bin\" / \"java\"), \"-version\"], check=True)\n\nensure_spark_java()\nos.environ[\"PYSPARK_PYTHON\"] = sys.executable\nfrom pyspark.sql import SparkSession\n_setup_spark = (SparkSession.builder.master(\"local[2]\").appName(\"cs145-p2\")\n    .config(\"spark.driver.memory\", \"2g\").getOrCreate())\ntry:\n    assert _setup_spark.range(1).count() == 1\n    print(\"PASS: Java and Spark setup\")\nfinally:\n    _setup_spark.stop()\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-002"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-003"
      },
      "source": "### Try Spark on a small example\n\nSpark selects relevant documents and removes duplicates. Run these three cells, then change `keyword` and rerun the last cell. This practice data is separate from your project collection; your project tables stay in BigQuery.\n",
      "id": "science-v2-003"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-004"
      },
      "source": "# 1. Create a tiny practice collection.\nfrom pyspark.sql import SparkSession, functions as F\npractice_spark = (SparkSession.builder.master(\"local[2]\").appName(\"cs145-p2\")\n    .config(\"spark.driver.memory\", \"2g\").getOrCreate())\npractice_docs = practice_spark.createDataFrame([\n    (\"1\", \"Urban rail\", \"Trains carry passengers between city stations.\"),\n    (\"2\", \"Airport travel\", \"Trains and buses connect the airport to the city.\"),\n    (\"3\", \"Bicycle routes\", \"Protected lanes help people cycle across the city.\"),\n    (\"2\", \"Airport travel\", \"Trains and buses connect the airport to the city.\"),\n], [\"id\", \"title\", \"text\"])\npractice_docs.select(\"id\", \"title\").show(truncate=False)\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-004"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-005"
      },
      "source": "# 2. Remove the repeated document.\npractice_unique = practice_docs.dropDuplicates([\"id\"])\npractice_unique.orderBy(\"id\").show(truncate=False)\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-005"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-006"
      },
      "source": "# 3. Change \"trains\" to \"cycle\" and rerun.\nkeyword = \"trains\"\npractice_unique.filter(F.lower(F.col(\"text\")).contains(keyword.lower())).orderBy(\"id\").show(truncate=False)\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-006"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-007"
      },
      "source": "### Optional: save checkpoints in Drive\n\nIn Colab, uncomment and run the cell below before collecting documents. It saves files in your Drive across runtime resets. On reconnect, rerun Setup and this cell to reuse them. Saving the notebook alone does not save local files.\n",
      "id": "science-v2-007"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-008"
      },
      "outputs": [],
      "execution_count": null,
      "source": "# from google.colab import drive\n# drive.mount('/content/drive')\n# WORK = Path('/content/drive/MyDrive/cs145_nanoquery/p2_work')\n# WORK.mkdir(parents=True, exist_ok=True)\n",
      "id": "science-v2-008"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-009"
      },
      "source": "# @title SQL and retrieval helpers\nimport sqlglot\nfrom sqlglot import exp\n# Restrict queries to SELECT-family statements. Pair with read-only IAM in production.\ndef check_select(sql):\n    statements = sqlglot.parse(sql, read=\"bigquery\")\n    if len(statements) != 1 or not isinstance(statements[0], (exp.Select, exp.Union, exp.Intersect, exp.Except)):\n        raise ValueError(\"Expected one SELECT query\")\n    return sql\n\ndef run_select(sql, parameters=None):\n    assert RUN_BIGQUERY, \"Enable BigQuery after setup\"\n    check_select(sql)\n    preview = preview_query(sql, parameters)\n    if not preview[\"within_limit\"]:\n        raise ValueError(\"Estimated scan exceeds MAX_BYTES_BILLED. Narrow the scan or deliberately change the limit; it is never raised automatically.\")\n    config = bigquery.QueryJobConfig(query_parameters=parameters or [],\n                                    maximum_bytes_billed=MAX_BYTES_BILLED,\n                                    use_query_cache=False)\n    start = time.perf_counter()\n    job = bq.query(sql, job_config=config)\n    rows = [dict(row) for row in job.result()]\n    return {\"rows\": rows, \"seconds\": time.perf_counter()-start,\n            \"bytes_processed\": job.total_bytes_processed}\n\ndef preview_query(sql, parameters=None):\n    assert RUN_BIGQUERY, \"Enable BigQuery after setup\"\n    check_select(sql)\n    config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False,\n                                    query_parameters=parameters or [])\n    job = bq.query(sql, job_config=config)\n    estimated = job.total_bytes_processed\n    if estimated is None:\n        raise ValueError(\"BigQuery did not provide a scan estimate; inspect the query before executing it\")\n    report = {\"estimated_bytes\": estimated, \"limit_bytes\": MAX_BYTES_BILLED,\n              \"within_limit\": estimated <= MAX_BYTES_BILLED}\n    print(report)\n    return report\n\ndef search_passages(question, k=5):\n    assert 1 <= k <= 50\n    vector = embedder.encode([question], normalize_embeddings=True)[0]\n    sql = f\"\"\"SELECT base.document_id, base.passage_id, base.title, base.url, base.text, distance\n    FROM VECTOR_SEARCH(TABLE `{PROJECT_ID}.{DATASET_ID}.wiki_passages`, 'embedding',\n      (SELECT @vector AS embedding), top_k => {k}, distance_type => 'COSINE',\n      options => '{{\"use_brute_force\":true}}')\n    ORDER BY distance\"\"\"\n    return run_select(sql, [bigquery.ArrayQueryParameter(\"vector\", \"FLOAT64\", vector.tolist())])\n\n# Run this cell as provided. No student work here for Phase 0.\n# Use search_passages() later, after preparing embeddings and the passage table.\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-009"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-010"
      },
      "source": "# @title Compare query results\nfrom collections import Counter\nfrom decimal import Decimal\nfrom numbers import Number\nfrom datetime import date, datetime\n\ndef canonical_cell(value):\n    if value is None:\n        return (\"null\",)\n    if isinstance(value, (float, np.floating)) and np.isnan(value):\n        return (\"null\",)\n    if isinstance(value, (bool, np.bool_)):\n        return (\"bool\", bool(value))\n    if isinstance(value, (Number, Decimal)):\n        return (\"number\", Decimal(str(value)).normalize())\n    if isinstance(value, (date, datetime)):\n        return (\"date\", value.isoformat())\n    return (\"text\", str(value))\n\ndef results_equal(actual, expected, columns, ordered=False):\n    # Pass run_select(...)[\"rows\"] and an explicit output-column order.\n    # Round approximate numeric results to your documented precision first.\n    from collections.abc import Mapping\n    if not columns or len(set(columns)) != len(columns):\n        raise ValueError(\"Supply unique output-column names in comparison order\")\n    def normalize(rows):\n        normalized = []\n        for row in rows:\n            if not isinstance(row, Mapping) or set(row) != set(columns):\n                raise ValueError(\"Each row must be a dictionary with the expected columns\")\n            normalized.append(tuple(canonical_cell(row[c]) for c in columns))\n        return normalized\n    a, e = normalize(actual), normalize(expected)\n    return a == e if ordered else Counter(a) == Counter(e)\n\n# Example: results_equal(actual[\"rows\"], expected[\"rows\"], columns=[\"zone\", \"trips\"])\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-010"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-011"
      },
      "source": "# @title Read pasted SQL\ndef strip_model_fence(raw):\n    text = raw.strip()\n    match = re.fullmatch(r\"```(?:json|sql)?\\s*\\n?(.*?)\\n?```\", text, re.DOTALL | re.IGNORECASE)\n    return match.group(1).strip() if match else text\n\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-011"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-012"
      },
      "source": "# @title Build prompts and save responses\ndef build_ab_messages(question, schema_notes, library, *, candidate_ids=(),\n                      rules=\"\", output_columns=()):\n    \"\"\"Build fixed prompts; keep reference SQL and answer values out of the inputs.\"\"\"\n    if not library:\n        raise ValueError(\"Populate your SQL library before building prompts\")\n    common = {\"question\": question, \"schema\": schema_notes,\n        \"candidate_ids\": list(candidate_ids), \"output_columns\": list(output_columns)}\n    governed = {\"domain_rules\": rules, \"definitions\": library}\n    instruction = (\n        \"Write one BigQuery SELECT query for the question using only the supplied tables. \"\n        \"Return only SQL with the requested output-column names. \"\n        \"Use @candidate_ids when filtering by supplied entity IDs; the notebook binds them. \"\n        \"Use concrete values from the question for other filters. \"\n        \"When a tested SQL library and domain rules are provided, follow those definitions and adapt or combine \"\n        \"the SQL examples as needed.\")\n    return {\n        \"A\": [{\"role\": \"system\", \"content\": instruction},\n              {\"role\": \"user\", \"content\": json.dumps(common, ensure_ascii=False)}],\n        \"B\": [{\"role\": \"system\", \"content\": instruction},\n              {\"role\": \"user\", \"content\": json.dumps(\n                  {**common, \"governed_context\": governed}, ensure_ascii=False)}]}\n\n\ndef prompt_for_chat(messages):\n    return (chr(10) * 2).join(m[\"content\"] for m in messages)\n\ndef save_response(question_id, variant, response):\n    if variant not in (\"A\", \"B\") or not response.strip():\n        raise ValueError(\"Provide an A/B label and the original response\")\n    record = PASTED_RESPONSES.setdefault(question_id, {})\n    if variant in record:\n        raise ValueError(\"First response already saved; record reruns separately\")\n    record[variant] = response\n    RESPONSE_FILE.write_text(json.dumps(PASTED_RESPONSES, indent=2))\n\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-012"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-013"
      },
      "source": "# @title BigQuery connection helper\ndef connect_bigquery():\n    global bq, bigquery\n    assert re.fullmatch(r\"[a-z][a-z0-9-]{4,61}[a-z0-9]\", PROJECT_ID)\n    assert re.fullmatch(r\"[A-Za-z_][A-Za-z0-9_]*\", DATASET_ID)\n    try:\n        import google.colab\n    except ImportError:\n        pass  # Use local Application Default Credentials.\n    else:\n        from google.colab import auth\n        auth.authenticate_user()\n    from google.cloud import bigquery\n    bq = bigquery.Client(project=PROJECT_ID, location=LOCATION)\n    dataset = bigquery.Dataset(f\"{PROJECT_ID}.{DATASET_ID}\")\n    dataset.location = LOCATION\n    bq.create_dataset(dataset, exists_ok=True)\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-013"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-014"
      },
      "source": "## Configuration\n",
      "id": "science-v2-014"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-015"
      },
      "source": "Use your own **Google Cloud project ID** from the console's project selector for `PROJECT_ID`, not the public dataset's project. `DATASET_ID` names the output dataset this notebook creates in your project. Set `LOCATION` to your source dataset's location, and `SOURCE_TABLE` to an existing `project.dataset.table` ID you can read. Your account needs permission to run query jobs and create the output dataset in your project.\n\nSet `RUN_BIGQUERY = True` to connect. Set `RUN_COLLECTION = True` after filling in `TITLES` and `WIKIPEDIA_CONTACT` in Phase 0. Sign in with the account that has access to your project. BigQuery sandbox is sufficient within its quotas; no paid LLM API is required.\n\n`MAX_BYTES_BILLED` defaults to 100 MB **per query**, not per dataset. `preview_query(sql)` prints the estimated scan and configured limit; `run_select` checks it before executing. Narrow an over-limit scan or deliberately change the limit within your available quota. The helper never raises it automatically. `LIMIT` controls returned rows and generally does not reduce the scan.\n\nThe connection cell checks query access, reads one source row, and loads and reads a one-row temporary setup table. It removes that setup table afterward. All three checks must pass before continuing.\n",
      "id": "science-v2-015"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-016"
      },
      "source": "RUN_COLLECTION = False\nRUN_BIGQUERY = False\nPROJECT_ID = \"\"  # Your project; do not enter billing credentials here.\nDATASET_ID = \"cs145_p2\"\nLOCATION = \"US\"\nSOURCE_TABLE = \"\"  # One existing source table: project.dataset.table\n# This 100 MB limit applies to each query scan, not total dataset size (500 MB+).\n# Call preview_query(sql) to see the estimated bytes before running a query.\n# If needed, reduce the scan or deliberately raise this cap within sandbox limits.\nMAX_BYTES_BILLED = 100_000_000\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-016"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-017"
      },
      "source": "# @title Connect to BigQuery (also used for the proposal)\nif RUN_BIGQUERY:\n    connect_bigquery()\n    assert re.fullmatch(r\"[A-Za-z0-9_-]+\\.[A-Za-z0-9_]+\\.[A-Za-z0-9_-]+\", SOURCE_TABLE), \"Fill SOURCE_TABLE with project.dataset.table\"\n    assert run_select(\"SELECT 1 AS connected\")[\"rows\"] == [{\"connected\": 1}]\n    print(\"PASS: BigQuery query access\")\n    source = bq.get_table(SOURCE_TABLE)\n    if source.location.lower() != LOCATION.lower():\n        raise ValueError(f\"Source is in {source.location}; set LOCATION to match before continuing\")\n    sample_rows = list(bq.list_rows(source, max_results=1))\n    print(\"PASS: source table access\", SOURCE_TABLE, \"sample rows:\", len(sample_rows))\n    from uuid import uuid4\n    setup_table = f\"{PROJECT_ID}.{DATASET_ID}._setup_check_{uuid4().hex}\"\n    try:\n        config = bigquery.LoadJobConfig(schema=[bigquery.SchemaField(\"connected\", \"INTEGER\")])\n        bq.load_table_from_json([{\"connected\": 1}], setup_table, job_config=config).result()\n        assert run_select(f\"SELECT connected FROM `{setup_table}`\")[\"rows\"] == [{\"connected\": 1}]\n        print(\"PASS: output dataset load and read\")\n    finally:\n        bq.delete_table(setup_table, not_found_ok=True)\n    print(\"Setup complete. Continue to Phase 0.\")\nelse:\n    print(\"Fill Configuration and set RUN_BIGQUERY = True, then rerun this cell.\")\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-017"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "phase-0"
      },
      "source": "## Phase 0: Prepare documents\n\n1. **Choose your tables:** Use your Project 1 dataset. You can also choose a different public dataset if you’d like; the same Project 1 data requirements apply.\n2. **Add documents:** Select 100–200 related Wikipedia articles or supply custom domain documents, prepare them with Spark, and map them to your table records. Inspect a random sample of 30 documents for relevance (all if fewer than 30).\n3. **Query and submit:** Run one keyword join combining documents and tables, create one chart, and complete the [proposal checklist](https://cs145.web.app/projects/project2-science.html#proposal).\n\nEmbeddings and LLM runs are not required for the proposal. For custom documents, supply an extract with id, url (or a stable source path), title, and text; use the same Spark preparation, mapping, and evidence checks.\n\n<a id=\"scrollTo=phase-0\"></a>\n",
      "id": "phase-0"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-019"
      },
      "source": "DOCUMENT_SOURCE = \"wikipedia\"  # Use \"custom\" for your own document extract.\nTITLES = []  # Fill with 100–200 Wikipedia article titles; fewer for development.\nMY_KEYWORDS = []  # Domain keywords used in your Spark filter.\nWIKIPEDIA_CONTACT = \"\"  # A contact URL or email for the public API User-Agent.\nEXTRACT_PATH = WORK / \"wiki_extract.parquet\"\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-019"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-020"
      },
      "source": "# @title Download selected Wikipedia articles\nif RUN_COLLECTION:\n    import requests\n    from datetime import datetime, timezone\n    assert WIKIPEDIA_CONTACT and 1 <= len(TITLES) <= 200\n    session = requests.Session()\n    session.headers[\"User-Agent\"] = f\"CS145CourseProject/1.0 ({WIKIPEDIA_CONTACT})\"\n    rows, missing, duplicates = [], [], []\n    seen_titles, seen_ids = set(), set()\n    for raw_title in TITLES:\n        if not isinstance(raw_title, str) or not raw_title.strip():\n            raise ValueError(\"Use nonempty article titles in TITLES\")\n        title = raw_title.strip()\n        if title in seen_titles:\n            duplicates.append({\"requested\": title, \"reason\": \"repeated title\"})\n            continue\n        seen_titles.add(title)\n        response = session.get(\"https://en.wikipedia.org/w/api.php\", params={\n            \"action\": \"query\", \"format\": \"json\", \"formatversion\": 2,\n            \"prop\": \"extracts|info\", \"explaintext\": 1, \"inprop\": \"url\",\n            \"redirects\": 1, \"titles\": title}, timeout=60)\n        response.raise_for_status()\n        payload = response.json()\n        if \"error\" in payload:\n            raise RuntimeError(payload[\"error\"])\n        pages = payload[\"query\"][\"pages\"]\n        usable = [p for p in pages if not p.get(\"missing\") and p.get(\"extract\")]\n        if not usable:\n            missing.append(title)\n        for p in usable:\n            page_id = str(p[\"pageid\"])\n            if page_id in seen_ids:\n                duplicates.append({\"requested\": title, \"resolved\": p[\"title\"],\n                                   \"reason\": \"same article ID\"})\n                continue\n            seen_ids.add(page_id)\n            rows.append({\"id\": page_id, \"url\": p[\"fullurl\"],\n                         \"title\": p[\"title\"], \"text\": p[\"extract\"],\n                         \"collected_at\": datetime.now(timezone.utc).isoformat()})\n        time.sleep(0.2)\n    report = {\"requested\": len(TITLES), \"downloaded_unique\": len(rows),\n              \"missing_or_empty\": missing, \"duplicates\": duplicates}\n    (WORK / \"collection_report.json\").write_text(json.dumps(report, indent=2))\n    print(json.dumps(report, indent=2))\n    if not rows:\n        raise ValueError(\"No articles downloaded. Check collection_report.json; existing extract unchanged.\")\n    pd.DataFrame(rows).to_parquet(EXTRACT_PATH, index=False)\n    print(\"Saved\", len(rows), \"unique articles to\", EXTRACT_PATH)\n    if len(rows) < 100:\n        print(\"Development collection: expand to 100–200 relevant articles for submission.\")\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-020"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-021"
      },
      "source": "# @title Prepare documents with Spark\n# Run after supplying a bounded extract and choosing keywords.\nif EXTRACT_PATH.exists() and MY_KEYWORDS:\n    assert DOCUMENT_SOURCE in (\"wikipedia\", \"custom\")\n    from pyspark.sql import SparkSession, functions as F\n    spark = SparkSession.builder.master(\"local[2]\").appName(\"cs145-p2\").config(\"spark.driver.memory\", \"2g\").getOrCreate()\n    pattern = \"|\".join(re.escape(k.lower()) for k in MY_KEYWORDS)\n    selected = (spark.read.parquet(str(EXTRACT_PATH))\n        .filter(F.col(\"text\").isNotNull())\n        .filter(F.lower(F.concat_ws(\" \", \"title\", \"text\")).rlike(pattern))\n        .dropDuplicates([\"id\"]).orderBy(\"id\"))\n    articles = selected.toPandas()\n    assert len(articles), \"No matching articles; inspect your filter\"\n    articles.to_parquet(WORK / \"articles.parquet\", index=False)\n    display(articles[[\"title\", \"text\"]].head(3).assign(text=lambda df: df.text.str.slice(0, 250)))\n    sample = articles.sample(n=min(30, len(articles)), random_state=145)\n    relevance_path = WORK / \"article_relevance.csv\"\n    review = sample.assign(relevant=\"\", explanation=\"\")\n    if relevance_path.exists():\n        previous = pd.read_csv(relevance_path, dtype={\"id\": str}, keep_default_na=False)\n        # Keep a backup, including labels for documents outside the current sample.\n        from uuid import uuid4\n        relevance_path.with_name(f\"article_relevance_backup_{uuid4().hex}.csv\").write_bytes(relevance_path.read_bytes())\n        previous = previous.drop_duplicates(\"id\", keep=\"last\").set_index(\"id\")\n        for column in [\"relevant\", \"explanation\"]:\n            review[column] = review[\"id\"].astype(str).map(previous[column]).fillna(\"\")\n    review.to_csv(relevance_path, index=False)\n    if DOCUMENT_SOURCE == \"wikipedia\" and not 100 <= len(articles) <= 200:\n        print(\"Development collection: keep 100–200 relevant Wikipedia articles after filtering for submission.\")\n    print(\"Review labels:\", relevance_path, \"(existing labels preserved by document ID)\")\n    print(\"Selected articles:\", len(articles), \"UTF-8 text bytes:\", sum(len(t.encode()) for t in articles.text))\n    spark.stop()\n    # Release the local JVM as well as the Spark session before computing embeddings.\n    from pyspark import SparkContext\n    if SparkContext._gateway is not None:\n        SparkContext._gateway.shutdown()\n    SparkContext._gateway = None\n    SparkContext._jvm = None\n    print(\"Spark stopped. Restart the kernel before starting another Spark session.\")\nelse:\n    print(\"Supply wiki_extract.parquet and MY_KEYWORDS, then run this cell.\")\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-021"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-022"
      },
      "source": "### Try your documents\n\nRun the supplied filter, then enter two keywords or phrases relevant to your project and rerun. Keep the outputs visible. Which documents appear or disappear? No extra paragraph is needed.\n\nThese are keyword matches in your prepared documents, ordered by document ID—not relevance rankings. This reads Spark’s saved output without changing it.\n",
      "id": "science-v2-022"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-023"
      },
      "source": "# @title Document filter helpers\nfrom html import escape\nfrom IPython.display import HTML, display\n\ndef show_document_matches(documents, phrase, limit=3):\n    phrase = phrase.strip()\n    if not phrase:\n        return\n    needle = phrase.lower()\n    ordered = documents.assign(id=documents.id.astype(str)).sort_values(\"id\", kind=\"stable\")\n    matches = ordered[\n        ordered.title.fillna(\"\").str.lower().str.contains(needle, regex=False) |\n        ordered.text.fillna(\"\").str.lower().str.contains(needle, regex=False)]\n    print(f\"Filter: {phrase} — {len(matches)} matching documents. Showing up to {limit}.\")\n    rows = []\n    for row in matches.head(limit).itertuples():\n        text = str(row.text)\n        match = re.search(re.escape(phrase), text, flags=re.IGNORECASE)\n        start = max(0, match.start() - 70) if match else 0\n        end = max(start + 250, match.end()) if match else 250\n        excerpt = (\"…\" if start else \"\") + text[start:end] + (\"…\" if len(text) > end else \"\")\n        excerpt = re.sub(re.escape(phrase), lambda m: \"<strong>\" + m.group(0) + \"</strong>\", escape(excerpt), flags=re.IGNORECASE) if not any(ch in phrase for ch in \"<>&\") else escape(excerpt)\n        rows.append(f\"<tr><td>{escape(str(row.id))}</td><td>{escape(str(row.title))}</td><td>{excerpt}</td></tr>\")\n    if rows:\n        display(HTML(\"<table><tr><th>ID</th><th>Document</th><th>Excerpt</th></tr>\" + \"\".join(rows) + \"</table>\"))\n    else:\n        print(\"No matches. Try a different keyword or phrase.\")\n    return matches.id.tolist()\n\nprepared_path = WORK / \"articles.parquet\"\nif not prepared_path.exists():\n    raise FileNotFoundError(\"Run Prepare documents with Spark first.\")\nquery_documents = pd.read_parquet(prepared_path)\n# Choose a keyword that actually occurs in this collection.\nsearchable = query_documents.title.fillna(\"\") + \" \" + query_documents.text.fillna(\"\")\nexample_filter = next((k for k in MY_KEYWORDS if k.strip() and searchable.str.contains(k, case=False, regex=False).any()), None)\nif example_filter is None:\n    example_filter = next((word for text in searchable for word in re.findall(r\"\\w+\", text) if len(word) > 2), None)\nif example_filter is None:\n    raise ValueError(\"No example keyword found; inspect your prepared documents.\")\nprint(f\"Prepared collection: {len(query_documents)} documents. Example filter: {example_filter}\")\ndisplay(query_documents[[\"id\", \"title\", \"text\"]].head(5).assign(text=lambda df: df.text.str.slice(0, 180)))\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-023"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-024"
      },
      "source": "filters = [\n    example_filter,  # Filter 1: a keyword from your collection.\n    \"\",              # Filter 2: enter your own keyword or phrase.\n    \"\",              # Filter 3: try another.\n]\nfor phrase in filters:\n    if phrase.strip():\n        _ = show_document_matches(query_documents, phrase)\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-024"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-025"
      },
      "source": "# STUDENT WORK: label the saved CSV, then run this cell to show your labels\n# and fraction relevant. Keep the output visible; no separate CSV is needed.\n# Use yes/no (or 1/0) in the relevant column.\nrelevance_path = WORK / \"article_relevance.csv\"\nif relevance_path.exists():\n    labeled = pd.read_csv(relevance_path, keep_default_na=False)\n    display(labeled[[\"title\", \"relevant\", \"explanation\"]])\n    labels = labeled[\"relevant\"].astype(str).str.strip().str.lower().map(\n        {\"yes\": 1, \"no\": 0, \"1\": 1, \"0\": 0, \"true\": 1, \"false\": 0})\n    if len(labels) and labels.notna().all():\n        print(f\"Fraction relevant: {int(labels.sum())}/{len(labels)} = {labels.mean():.1%}\")\n    else:\n        print(\"Label every sampled document in the CSV, save it, and rerun this cell.\")\nelse:\n    print(\"Run the document-preparation cell first, then label its CSV.\")\n\n# Map documents to table IDs, then run one keyword join with a chart.\n# Keep these outputs visible too.\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-025"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-026"
      },
      "source": "### Proposal summary\n\nFill in the answers below. Refer to your outputs above rather than repeating them.\n\n| Question | Your answer |\n|---|---|\n| Who is on your team, and how will you divide the work? | [Answer] |\n| Which tier are you targeting? | [Answer] |\n| Which tables and document sources will you use? | [Answer] |\n| What analytical questions will you investigate? | [Answer] |\n",
      "id": "science-v2-026"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "phase-1"
      },
      "source": "**Proposal: stop here. Continue to Phase 1 after completing Phase 0.**\n\n## Phase 1: Embeddings and semantic search\n\n<a id=\"scrollTo=phase-1\"></a>\n",
      "id": "phase-1"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-028"
      },
      "source": "# @title Compute MiniLM embeddings\nfrom sentence_transformers import SentenceTransformer\nEMBED_MODEL = \"sentence-transformers/all-MiniLM-L6-v2\"\nimport torch\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Embedding device:\", DEVICE)\nembedder = SentenceTransformer(EMBED_MODEL, device=DEVICE)\ndef make_passages(articles, embedder):\n    tok = embedder.tokenizer\n    size = embedder.max_seq_length - tok.num_special_tokens_to_add(pair=False) - 8\n    assert size > 0\n    rows = []\n    for article in articles.sort_values(\"id\").to_dict(\"records\"):\n        tokens = tok.encode(article[\"text\"], add_special_tokens=False)\n        offset = 0\n        while offset < len(tokens):\n            count = min(size, len(tokens) - offset)\n            text = tok.decode(tokens[offset:offset+count], skip_special_tokens=True)\n            # Keep all source tokens: move unused tokens to the next passage.\n            while len(tok.encode(text, add_special_tokens=True)) > embedder.max_seq_length:\n                count -= 1\n                if count <= 0:\n                    raise ValueError(\"A source token cannot fit in the embedding context\")\n                text = tok.decode(tokens[offset:offset+count], skip_special_tokens=True)\n            if text.strip():\n                rows.append({\"document_id\": str(article[\"id\"]),\n                    \"passage_id\": f'{article[\"id\"]}:{offset}', \"title\": article[\"title\"],\n                    \"url\": article[\"url\"], \"text\": text})\n            offset += count\n    return pd.DataFrame(rows)\nif (WORK / \"articles.parquet\").exists():\n    passages = make_passages(pd.read_parquet(WORK / \"articles.parquet\"), embedder)\n    print(f\"Embedding {len(passages):,} passages.\")\n    if len(passages) > 20_000:\n        print(\"Large collection: embedding may take longer. Relevant sections can reduce runtime.\")\n    t0 = time.perf_counter()\n    vectors = embedder.encode(passages.text.tolist(), batch_size=32, normalize_embeddings=True)\n    assert vectors.shape == (len(passages), 384) and np.isfinite(vectors).all()\n    assert np.allclose(np.linalg.norm(vectors, axis=1), 1, atol=1e-5)\n    passages[\"embedding\"] = vectors.tolist()\n    passages.to_parquet(WORK / \"passages.parquet\", index=False)\n    print(\"Passages:\", len(passages), \"Embedding seconds:\", time.perf_counter()-t0)\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-028"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-029"
      },
      "source": "### BigQuery batch load\n",
      "id": "science-v2-029"
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "science-v2-030"
      },
      "source": "# @title Load embeddings into BigQuery\nif RUN_BIGQUERY:\n    passages = pd.read_parquet(WORK / \"passages.parquet\")\n    rows = [{**r, \"embedding\": [float(x) for x in r[\"embedding\"]]} for r in passages.to_dict(\"records\")]\n    schema = [bigquery.SchemaField(n, \"STRING\") for n in [\"document_id\", \"passage_id\", \"title\", \"url\", \"text\"]]\n    schema += [bigquery.SchemaField(\"embedding\", \"FLOAT64\", mode=\"REPEATED\")]\n    cfg = bigquery.LoadJobConfig(schema=schema, write_disposition=\"WRITE_TRUNCATE\")\n    bq.load_table_from_json(rows, f\"{PROJECT_ID}.{DATASET_ID}.wiki_passages\", job_config=cfg).result()\n    print(\"Loaded\", len(rows), \"passages\")\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-030"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-031"
      },
      "source": "## Section 1: Analytical SQL and charts\n\nComplete five analytical queries and three charts in total. Reuse your proposal query and chart if they meet these requirements.\n\nEach query must combine at least two of: joins, CTEs, window functions, subqueries, CASE-based analysis, or semantic retrieval. CASE-based analysis applies a classification or calculation used in the answer. Semantic retrieval followed by a join counts as two techniques.\n\nGive each chart a title, labeled axes, and a short caption. Keep your checks and results visible.\n",
      "id": "science-v2-031"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-032"
      },
      "source": "ANALYTICAL_QUERIES = []\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-032"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-033"
      },
      "source": "## Section 2: Documents and tables\n",
      "id": "science-v2-033"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-034"
      },
      "source": "# STUDENT WORK: use semantic retrieval for two of your five analytical queries.\n# Save excerpts/source URLs, mappings, SQL/results, and checks for each.\n# Map retrieved document IDs to IDs in your tables; bind those entity IDs as SQL parameters.\nDOCUMENT_QUERY_CHECKS = []\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-034"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-035"
      },
      "source": "## Section 3: Query analysis\n",
      "id": "science-v2-035"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-036"
      },
      "source": "# STUDENT WORK: choose one query and record its measured bytes processed.\n# State operator/page/buffer assumptions; estimate I/O at current, 10x,\n# and 100x sizes. The larger sizes are estimates, not required runs.\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-036"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-037"
      },
      "source": "## Section 4: Findings and AI takeaways\n",
      "id": "science-v2-037"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-038"
      },
      "source": "Write half a page to one page maximum: findings, 2–3 examples of AI assistance and how you checked them, and a plain-English walkthrough of one part of the pipeline. Code, figures, tables, and Section 7 evaluation are outside this prose limit. If you did not use AI assistance, say so and describe 2–3 checks, debugging decisions, or revisions you made independently. Tier 100 still requires the LLM experiment.\n\n[Your writeup]\n",
      "id": "science-v2-038"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-039"
      },
      "source": "## Section 5: Governed SQL library\n",
      "id": "science-v2-039"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-040"
      },
      "source": "LIBRARY = {}  # Populate five entries.\n# Shape of an entry (write your own SQL):\n# LIBRARY[\"definition_name\"] = {\n#   \"description\": \"...\", \"domain_rules\": \"...\", \"parameters\": {\"budget\": \"FLOAT64\"},\n#   \"output_columns\": [\"...\"], \"sql\": \"SELECT ... WHERE price < @budget\"}\n\n# STUDENT WORK: definitions, manifest, reference notes, tests, and candidate-ID binding.\n# One small fixture per definition; cover NULLs, duplicates, and missing matches.\nLIBRARY_TESTS = []\n# LIBRARY_TESTS.append({\n#     \"definition\": \"definition_name\",\n#     \"input_rows\": {\"table_name\": []},  # Fill with a tiny, explicit dataset.\n#     \"parameters\": {},\n#     \"expected_rows\": [],  # Compute by hand before running the SQL.\n#     \"output_columns\": [],\n# })\n# Run the definition's SQL against these fixture rows, then check:\n# assert results_equal(actual_rows, test[\"expected_rows\"], test[\"output_columns\"])\n# Use inline CTEs or separate fixture tables so tests leave analytical data intact.\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-040"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-041"
      },
      "source": "## Section 6: Ground truth (Tier 93)\n\nSave the LIBRARY version used for the experiment and keep it fixed across the A/B comparisons. Choose five new questions that reuse or combine its definitions. Include two document/table questions and two covering NULLs, duplicates, or missing matches (they may overlap). Write and run reference SQL; save the expected rows and explain your checks. No A/B calls are required here.",
      "id": "science-v2-041"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-042"
      },
      "source": "# Five new questions and verified reference answers belong to Tier 93.\nEVALUATION_QUESTIONS = []\n# Entry: {\"id\": \"q1\", \"question\": \"...\", \"reference_sql\": \"...\",\n#         \"reference_rows\": [], \"output_columns\": [], \"candidate_ids\": []}\n# Keep reference SQL and rows out of BOTH prompts; reuse these records in Section 7.\n# STUDENT WORK: run each reference query, save rows, and document independent checks.\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-042"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-043"
      },
      "source": "## Section 7: Text-to-SQL comparison and evaluation (Tier 100)\n\nUse the same five questions and frozen library. A receives base context; B also receives LIBRARY. Keep reference SQL and rows out of both prompts. Use the same model/settings and save the ten original responses. Apply identical safety checks, run the SQL, and compare results with your Tier 93 answers. Report correct answers out of five for each version and explain errors.",
      "id": "science-v2-043"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-044"
      },
      "source": "# After filling LIBRARY and preparing a question, use:\n# messages = build_ab_messages(question, schema_notes, LIBRARY,\n#     candidate_ids=ids, rules=domain_rules, output_columns=columns)\n# Print each prompt and paste it into a fresh LLM chat:\n# print(prompt_for_chat(messages[\"A\"]))\n# print(prompt_for_chat(messages[\"B\"]))\n\n# Paste the original responses into this dictionary in the submitted notebook.\n# Fill it before rerunning; this cell restores responses even in a fresh runtime.\nSUBMITTED_RESPONSES = {}  # {\"q1\": {\"A\": \"SELECT ...\", \"B\": \"SELECT ...\"}, ...}\nRESPONSE_FILE = WORK / \"llm_responses.json\"\nPASTED_RESPONSES = json.loads(RESPONSE_FILE.read_text()) if RESPONSE_FILE.exists() else {}\nfor question_id, responses in SUBMITTED_RESPONSES.items():\n    saved = PASTED_RESPONSES.setdefault(question_id, {})\n    for variant, response in responses.items():\n        if variant in saved and saved[variant] != response:\n            raise ValueError(\"Submitted response differs from saved first response; inspect before continuing\")\n        saved[variant] = response\n\n# STUDENT WORK: prepare prompts with build_ab_messages(); save pasted responses.\n# For BOTH A and B: strip_model_fence(), then run_select() with identical checks.\n# Add a table allowlist check before executing either response.\n# List allowed project.dataset.table IDs. Check physical tables in joins,\n# CTE bodies, and subqueries; CTE names and aliases are not physical tables.\n# Reject unapproved or unresolved references before calling run_select().\n# Demonstrate one allowed query and one rejection. Save original SQL and the\n# rejection reason, skip execution, and continue. Do not silently repair SQL.\n# Return source IDs/URLs, SQL, results, and original outputs. Catch and log failures.\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-044"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "science-v2-045"
      },
      "source": "# Reuse EVALUATION_QUESTIONS and the A/B results collected above.\n# STUDENT WORK: compare results with results_equal(), preserve duplicates,\n# check ordering when requested, document numeric precision, and explain failures.\n",
      "outputs": [],
      "execution_count": null,
      "id": "science-v2-045"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "science-v2-046"
      },
      "source": "## Submitting to Gradescope\n\nComplete through Section 4 for T80, Section 5 for T87, Section 6 for T93, or Section 7 for T100. Submit each section’s results once; mark later sections “Not attempted.”\n\n- Submit your final notebook(s) as both `.ipynb` and PDF files, with saved outputs visible.\n- Tag the PDF pages to the associated Gradescope questions.\n- In Colab, choose **File → Print**, then **Save as PDF**.\n",
      "id": "science-v2-046"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}