notebooks/gemini_labelling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c206dfbb-2223-43e4-9f0c-f0b243273f70",
   "metadata": {},
   "source": [
    "# Sentence-Query Pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a380067-0872-4d77-bcbd-627ff68e6d92",
   "metadata": {},
   "source": [
    "## -> Create S-Q Pairs Dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab097245-eade-4478-8ac9-6cb2d4480cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "create_new_dictionary = False\n",
    "\n",
    "if create_new_dictionary:\n",
    "    from utilities import load_jsonl\n",
    "    \n",
    "    # identify datasets to use\n",
    "    dataset_names = ['fava', 'wiki']\n",
    "    \n",
    "    # identify query outputs to look for\n",
    "    query_experiments = ['q1', 'q2', 'q3', 'q4', 'q5', 'q6', 'q7']\n",
    "    \n",
    "    \n",
    "    # final result will be a dictionary\n",
    "    result = {\"sentence_query_relevance\": {}, \"included_queries\": {\"fava\": [], \"wiki\": []}}\n",
    "    \n",
    "    # load fava and wiki outputs \n",
    "    for ds_name in dataset_names:\n",
    "        for qe in query_experiments:\n",
    "            file_path = f\"/Users/jjr/output/rarr-rep/output/{ds_name}/{ds_name}_{qe}_r1_a1.jsonl\"\n",
    "            try:\n",
    "                data = load_jsonl(file_path)\n",
    "                result[\"included_queries\"][ds_name].append(qe)\n",
    "            except FileNotFoundError:\n",
    "                print(f\"no experiment at {file_path}\")\n",
    "                # do not add this datasets s-e pairs\n",
    "                continue\n",
    "            \n",
    "            # iterate over sentences\n",
    "            for i in range(len(data)):\n",
    "                sentence = data[i][\"decon_sentence\"]\n",
    "        \n",
    "                # iterate over queries \n",
    "                queries = data[i].get('queries', [])\n",
    "                for query in queries:\n",
    "        \n",
    "                    # check if s-q pair is already in dictionary\n",
    "                    if (sentence, query) in result[\"sentence_query_relevance\"]:\n",
    "                        # skip to next query   \n",
    "                        continue\n",
    "                    # add s-q pair\n",
    "                    else:\n",
    "                        result[\"sentence_query_relevance\"][(sentence, query)] = None\n",
    "                # if wiki also check facts\n",
    "                if ds_name == 'wiki':\n",
    "                    facts = data[i].get('facts', [])\n",
    "                    for fact in queries:\n",
    "                        # check if s-q pair is already in dictionary\n",
    "                        if (sentence, fact) in result[\"sentence_query_relevance\"]:\n",
    "                            # skip to next query   \n",
    "                            continue\n",
    "                        # add s-q pair\n",
    "                        else:\n",
    "                            result[\"sentence_query_relevance\"][(sentence, fact)] = None\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    import pickle\n",
    "    \n",
    "    sq_dict_path = \"/Users/jjr/output/rarr-rep/gemini_labels/sentence-query/sentence_query_labels_dict.pkl\"\n",
    "    with open(sq_dict_path, \"wb\") as f:\n",
    "        pickle.dump(result, f)\n",
    "    \n",
    "    print(\"s-e pair dictionary saved as result.pkl\")\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1d923d2-cec3-4a41-9847-93ae8131ee14",
   "metadata": {},
   "outputs": [],
   "source": [
    "result[\"included_queries\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d82b293a-cb7a-4439-8f09-10d3bfb8b34b",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(result[\"sentence_query_relevance\"].keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8420803e-f0df-4a9d-88d4-2a89b0f9a36d",
   "metadata": {},
   "source": [
    "## -> create s-q pairs for testing Gemini prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89a09312-0318-48ed-a27a-3d64d249078f",
   "metadata": {},
   "outputs": [],
   "source": [
    "wiki_data = load_jsonl(f\"/Users/jjr/output/rarr-rep/output/wiki/wiki_q1_r1_a1.jsonl\")\n",
    "fava_data = load_jsonl(f\"/Users/jjr/output/rarr-rep/output/fava/fava_q1_r1_a1.jsonl\")\n",
    "\n",
    "\n",
    "# use wiki indices 0, 1, 4\n",
    "\n",
    "# import pickle\n",
    "# # load main results dictionary\n",
    "# with open(\"/Users/jjr/output/rarr-rep/gemini_labels/query_evidence_labels_dict.pkl\", \"rb\") as f:\n",
    "#     query_evidence_labels_dict = pickle.load(f)\n",
    "\n",
    "# list(q for (q,e) in query_evidence_labels_dict if query_evidence_labels_dict[(q,e)]['O'] == 2)[200:300]\n",
    "\n",
    "bad_sentences = list(s for (s,q) in list(result[\"sentence_query_relevance\"].keys()) if ((len(s) <= 25)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dbd4e40-5e18-442a-b502-bb462b281bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(bad_sentences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d897b71-47a0-4619-b05e-9b530c480acc",
   "metadata": {},
   "outputs": [],
   "source": [
    "set(bad_sentences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d4e1151-0a6e-47c2-ae91-68c266814804",
   "metadata": {},
   "outputs": [],
   "source": [
    "bad_sentences = [\"I'm happy to help!\", \"Certainly!\", \"Absolutely!\", \n",
    "             \"Sure!\", \"I hope that helps!\", \"I hope this helps!\",\n",
    "             \"I'm happy to help!\" \"Yes, that's correct!\", \"Hello!\"]\n",
    "\n",
    "bad_queries = list(q for (s,q) in list(result[\"sentence_query_relevance\"].keys()) if s in bad_sentences)\n",
    "\n",
    "len(set(bad_queries))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "842dd6bf-5004-43a3-9f3f-1edf34e4e3b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "fava_8b = load_jsonl(\"/Users/jjr/output/rarr-rep/output/fava/fava_q1_r1_a1.jsonl\")\n",
    "fava_70b = load_jsonl(\"/Users/jjr/output/rarr-rep/output/fava/fava_q2_r1_a1.jsonl\")\n",
    "\n",
    "queries_8b = []\n",
    "queries_70b = []\n",
    "\n",
    "for i in range(len(fava_8b)):\n",
    "    decon_sentence = fava_8b[i][\"decon_sentence\"]\n",
    "    if decon_sentence in bad_sentences:\n",
    "        queries_8b.append(fava_8b[i]['queries'])\n",
    "\n",
    "for i in range(len(fava_70b)):\n",
    "    decon_sentence = fava_8b[i][\"decon_sentence\"]\n",
    "    if decon_sentence in bad_sentences:\n",
    "        queries_70b.append(fava_70b[i]['queries'])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e14b0bd-a109-44b7-98ff-54f0685d69cc",
   "metadata": {},
   "source": [
    "## -> Run Labeller"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25afac9a-4eb2-45b3-82e3-d5105154f996",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_sq_pairs = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "975a27c1-8957-47cf-b453-956f1afa4129",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = '''\n",
    "You are a quality judge, evaluating both the sentence and the query.\n",
    "\n",
    "Instructions:\n",
    "\n",
    "1. You are given a single sentence and a single query.\n",
    "\n",
    "2. Evaluate the **sentence** first:\n",
    "  - A “good” sentence (score 1) contains enough factual or contextual detail to allow for meaningful queries.\n",
    "  - A “poor” sentence (score 0) has little or no factual content (e.g., “Bye!”, “Thank you.”) or is too vague to support a useful query.\n",
    "\n",
    "3. Next, evaluate the **query** with respect to the sentence:\n",
    "  - A “good” query (score 1) directly addresses or clarifies a fact, detail, or relevant aspect of the sentence in a way that would be useful for a Google search.\n",
    "  - A “poor” query (score 0) does not address any meaningful detail in the sentence or is clearly unrelated or unhelpful.\n",
    "\n",
    "4. **Output**:\n",
    "  - You must output a **valid JSON object** with two integer keys: `\"sentence_score\"` and `\"query_score\"`, each set to `0` or `1`.\n",
    "  - Example: {example_json}\n",
    "\n",
    "5. Output **only** the JSON and no extra text or explanation.\n",
    "\n",
    "Sentence: \"{sentence}\"\n",
    "Query: \"{query}\"\n",
    "Output:'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9032cd96-3c9d-476b-9191-359537a54914",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "import re\n",
    "\n",
    "import json\n",
    "\n",
    "import json\n",
    "\n",
    "def parse_json_scores(response_text: str) -> dict:\n",
    "    \"\"\"\n",
    "    Expects 'response_text' to contain exactly one JSON object of the form:\n",
    "      {\n",
    "        \"sentence_score\": 0 or 1,\n",
    "        \"query_score\": 0 or 1\n",
    "      }\n",
    "    Returns the parsed dictionary, for example:\n",
    "      {\n",
    "        \"sentence_score\": 1,\n",
    "        \"query_score\": 0\n",
    "      }\n",
    "\n",
    "    Raises ValueError if parsing fails or if the keys/values are invalid.\n",
    "    \"\"\"\n",
    "    response_text = response_text.strip()\n",
    "    \n",
    "    # Attempt to parse the entire text as JSON\n",
    "    try:\n",
    "        data = json.loads(response_text)\n",
    "    except json.JSONDecodeError as e:\n",
    "        raise ValueError(f\"Failed to parse JSON: {e}\")\n",
    "    \n",
    "    # Check for required keys\n",
    "    expected_keys = {\"sentence_score\", \"query_score\"}\n",
    "    if not all(key in data for key in expected_keys):\n",
    "        raise ValueError(\"JSON must contain 'sentence_score' and 'query_score' keys.\")\n",
    "\n",
    "    # Validate that both scores are either 0 or 1\n",
    "    if data[\"sentence_score\"] not in [0, 1] or data[\"query_score\"] not in [0, 1]:\n",
    "        raise ValueError(\"'sentence_score' and 'query_score' must each be 0 or 1.\")\n",
    "    \n",
    "    # Optional: Ensure there are no extra keys\n",
    "    if len(data) != 2:\n",
    "        raise ValueError(\"Unexpected additional keys in JSON.\")\n",
    "    \n",
    "    return data\n",
    "\n",
    "\n",
    "def exponential_backoff(attempt, base_delay=3, max_delay=60.0):\n",
    "    delay = base_delay * (2 ** attempt)  # exponential growth\n",
    "    delay = min(delay, max_delay)\n",
    "    jitter = random.uniform(0, delay * 0.15)  # add up to 15% jitter\n",
    "    return delay + jitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3aedb24-167f-453d-8551-f34712486bd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_sq_pairs:\n",
    "    import pickle\n",
    "    import os\n",
    "    import time\n",
    "    import datetime\n",
    "    # import vertexai\n",
    "    # from vertexai.generative_models import GenerativeModel\n",
    "    from tqdm import tqdm\n",
    "    from google import genai\n",
    "\n",
    "    result = {\"sentence_query_relevance\": {}, \"included_queries\": {\"fava\": [], \"wiki\": []}}\n",
    "    \n",
    "    # file paths for storing results\n",
    "    root_path = \"/Users/jjr/output/rarr-rep/gemini_labels/sentence-query/\"\n",
    "    sentence_query_labels_dict_path = os.path.join(root_path, \"sentence_query_labels_dict.pkl\")\n",
    "    failed_api_calls_dict_path = os.path.join(root_path, \"failed_api_calls_dict.pkl\")\n",
    "    \n",
    "    # load main results dictionary\n",
    "    with open(sentence_query_labels_dict_path, \"rb\") as f:\n",
    "        sentence_query_labels_dict = pickle.load(f)\n",
    "        \n",
    "    # attempt to load any previously failed API calls; if not, start fresh.\n",
    "    if os.path.exists(failed_api_calls_dict_path):\n",
    "        with open(failed_api_calls_dict_path, \"rb\") as f:\n",
    "            failed_api_calls_dict = pickle.load(f)\n",
    "    else:\n",
    "        failed_api_calls_dict = {}\n",
    "    \n",
    "    # create list of (sentence, query) tuples that have None as value\n",
    "    no_eval_list = [\n",
    "        (sentence, query) \n",
    "        for (sentence, query), val in sentence_query_labels_dict[\"sentence_query_relevance\"].items()\n",
    "        if val is None and (sentence, query) not in failed_api_calls_dict\n",
    "    ]\n",
    "\n",
    "    ###### obtain api key #####\n",
    "    api_key = \"\"\n",
    "    # instantiate model\n",
    "    client = genai.Client(api_key=api_key)\n",
    "    max_retries = 5\n",
    "    \n",
    "    # iterate over tuples in no_eval_list\n",
    "    i = 0\n",
    "    for (sentence, query) in tqdm(no_eval_list, desc=\"Processing pairs\"):\n",
    "    \n",
    "\n",
    "        # Populate the prompt with the required data\n",
    "        prompt_data = { \"sentence\": sentence, \"query\": query, \"example_json\": \"{\\\"sentence_score\\\": 1, \\\"query_score\\\": 0}\"}\n",
    "    \n",
    "        # create llm messages\n",
    "        messages = prompt.format(**prompt_data)\n",
    "        \n",
    "        # send to gemini\n",
    "        success = False\n",
    "        last_exception = None\n",
    "        for attempt in range(max_retries):\n",
    "            try:\n",
    "                response = client.models.generate_content(model=\"gemini-1.5-pro-002\", contents=messages)\n",
    "                response_text = response.text\n",
    "                success = True\n",
    "                time.sleep(7)\n",
    "                break\n",
    "            except Exception as e:\n",
    "                last_exception = e\n",
    "                if hasattr(e, 'response') and getattr(e.response, 'status_code', None) == 429:\n",
    "                    delay = exponential_backoff(attempt)\n",
    "                    print(f\"HTTP 429 received. Waiting {delay:.2f} seconds before retrying.\")\n",
    "                else:\n",
    "                    delay = exponential_backoff(attempt)\n",
    "                    print(f\"API call error: {e}. Retrying in {delay:.2f} seconds.\")\n",
    "                time.sleep(delay)\n",
    "    \n",
    "        # if the API call was not successful, flag and continue (do not update the main dictionary)\n",
    "        if not success:\n",
    "            error_msg = f\"API call failed after {max_retries} attempts: {last_exception}\"\n",
    "            failed_api_calls_dict[(sentence, query)] = error_msg\n",
    "            continue\n",
    "        \n",
    "        try:\n",
    "            parsed_dict = parse_json_scores(response_text)\n",
    "        except Exception as e:\n",
    "            print(f\"JSON parsing error for pair ({sentence}, {query}): {e}\")\n",
    "            failed_api_calls_dict[(sentence, query)] = f\"JSON parsing error: {e}\"\n",
    "            continue  \n",
    "        \n",
    "        # save label in dictionary\n",
    "        sentence_query_labels_dict[\"sentence_query_relevance\"][sentence, query] = parsed_dict\n",
    "        i += 1\n",
    "    \n",
    "        # save progress every 1000 iterations to avoid losing good API calls\n",
    "        if i % 100 == 0:\n",
    "            print(f\"saving iteration {i}\")\n",
    "            # save timestamped version\n",
    "            timestamp = datetime.datetime.now().strftime('%d_%m_%y_%H%M')\n",
    "            with open(os.path.join(root_path, f\"sentence_query_labels_dict_{timestamp}.pkl\"), \"wb\") as f:\n",
    "                pickle.dump(sentence_query_labels_dict, f)\n",
    "            with open(os.path.join(root_path, f\"failed_api_calls_dict_{timestamp}.pkl\"), \"wb\") as f:\n",
    "                pickle.dump(failed_api_calls_dict, f)\n",
    "            # update the fixed file names so future runs will load the latest errors\n",
    "            with open(sentence_query_labels_dict_path, \"wb\") as f:\n",
    "                pickle.dump(sentence_query_labels_dict, f)\n",
    "            with open(failed_api_calls_dict_path, \"wb\") as f:\n",
    "                pickle.dump(failed_api_calls_dict, f) \n",
    "    \n",
    "    # Save any remaining work after the loop completes to the fixed file paths\n",
    "    with open(sentence_query_labels_dict_path, \"wb\") as f:\n",
    "        pickle.dump(sentence_query_labels_dict, f)\n",
    "    with open(failed_api_calls_dict_path, \"wb\") as f:\n",
    "        pickle.dump(failed_api_calls_dict, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac02bf79-7517-4e45-89e4-0e231155ce55",
   "metadata": {},
   "source": [
    "# Query-Evidence Pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f99f77c-ef87-491b-9e66-ee9ee0b7eaa5",
   "metadata": {},
   "source": [
    "## -> create q,e dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e86fdcc8-341a-407d-a5df-a1ede6a508f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from utilities import load_jsonl\n",
    "import json\n",
    "from typing import List, Dict\n",
    "\n",
    "def load_jsonl(path:str)->List[Dict]:\n",
    "    \"\"\"loads a jsonlist file and returns a list of json objects\"\"\"\n",
    "    output = []\n",
    "    try:\n",
    "        with open(path, 'r') as f:\n",
    "            for line in f:\n",
    "                data = json.loads(line)\n",
    "                output.append(data)\n",
    "    except FileNotFoundError:\n",
    "        print(f\"Error: The file at {path} does not exist.\")\n",
    "        raise \n",
    "    return output\n",
    "\n",
    "# create dictionary of all used evidence with key: chunked_id, value: <text>\n",
    "\n",
    "# create filenames from which to find evidence\n",
    "query_experiments = [f\"q{i}\" for i in range(1,8)]\n",
    "\n",
    "fava_filenames = [f\"fava_{qe}_r1_a1.jsonl\" for qe in query_experiments if qe not in [\"q3\", \"q4\"]]\n",
    "fava_filenames.extend(['fava_q1_r2.jsonl', 'fava_q1_r3.jsonl'])\n",
    "\n",
    "wiki_filenames = [f\"wiki_{qe}_r1_a1.jsonl\" for qe in query_experiments if qe not in []]\n",
    "wiki_filenames.extend(['wiki_q1_r2.jsonl', 'wiki_q1_r3.jsonl'])\n",
    "\n",
    "evidence_dict = {}\n",
    "\n",
    "\n",
    "for file_name in fava_filenames + wiki_filenames:\n",
    "    if file_name.startswith(\"fava\"):\n",
    "        data = load_jsonl(f\"/Users/jjr/output/rarr-rep/output/fava/{file_name}\")\n",
    "    elif file_name.startswith(\"wiki\"):\n",
    "        data = load_jsonl(f\"/Users/jjr/output/rarr-rep/output/wiki/{file_name}\")\n",
    "\n",
    "    # iterate over jsonl and extract evidence\n",
    "    for i in range(len(data)):\n",
    "        retrieved_evidence = data[i].get(\"retrieved_evidence\", [])\n",
    "\n",
    "        # flatten the retrieved evidence\n",
    "        retrieved_evidence = [evidence for evidence_list in retrieved_evidence for evidence in evidence_list]\n",
    "\n",
    "        # iterate over evidence objects\n",
    "        for obj in retrieved_evidence:\n",
    "            title = obj['title']\n",
    "            text = obj['text']\n",
    "            chunk_id = obj[\"chunked_id\"]\n",
    "\n",
    "            if chunk_id not in evidence_dict:\n",
    "                evidence_dict[chunk_id] = title + \": \" + text\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28e3935c-937b-473b-89d8-1e6dbb7929b7",
   "metadata": {},
   "source": [
    "## -> create empty (query, chunk_id): label dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49a75380-fe91-41db-a3ef-232891c6d0a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_evidence_label_dict = {}\n",
    "\n",
    "# iterate over files again\n",
    "for file_name in fava_filenames + wiki_filenames:\n",
    "    if file_name.startswith(\"fava\"):\n",
    "        data = load_jsonl(f\"/Users/jjr/output/rarr-rep/output/fava/{file_name}\")\n",
    "    elif file_name.startswith(\"wiki\"):\n",
    "        data = load_jsonl(f\"/Users/jjr/output/rarr-rep/output/wiki/{file_name}\")\n",
    "\n",
    "    # iterate over jsonl and extract (query,evidence) pairs\n",
    "    for i in range(len(data)):\n",
    "        queries = data[i].get('queries', [])\n",
    "        retrieved_evidence_list = data[i].get(\"retrieved_evidence\", [])\n",
    "\n",
    "        # should be one evidence list for each query in the same order as the queries\n",
    "        if len(queries) != len(retrieved_evidence_list):\n",
    "            raise ValueError(\"query evidence mismatch!\")\n",
    "\n",
    "        for query, evidence_list in zip(queries, retrieved_evidence_list):\n",
    "            # iterate over evidence in list\n",
    "            for evidence in evidence_list:\n",
    "                chunked_id = evidence['chunked_id']\n",
    "                if (query, chunked_id) not in query_evidence_label_dict:\n",
    "                    query_evidence_label_dict[query, chunked_id] = None\n",
    "\n",
    "# save as pickle file\n",
    "\n",
    "import pickle\n",
    "from datetime import datetime\n",
    "\n",
    "timestamp = datetime.now().strftime('%d_%m_%y_%H%M')\n",
    "\n",
    "root_path = \"/Users/jjr/output/rarr-rep/gemini_labels/query-evidence/\"\n",
    "\n",
    "with open(root_path + f\"query_evidence_labels_dict_empty.pkl\", \"wb\") as f:\n",
    "    pickle.dump(query_evidence_label_dict, f)\n",
    "\n",
    "with open(root_path + f\"evidence_dict.pkl\", \"wb\") as f:\n",
    "    pickle.dump(evidence_dict, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c48e097f-3178-458e-9c47-84877a9c13a8",
   "metadata": {},
   "source": [
    "## -> Run Labeller"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1470a355-458e-455c-b2af-b4fbcb939869",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_qe_pairs = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e75a29c6-68ed-4ba5-8497-261563e60eb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = '''\n",
    "You are a search quality rater evaluating the relevance of web pages.  Given a query and a web page, you must provide a score on an integer scale of 0 to 2 with the following meanings:\n",
    "2 = highly relevant, very helpful for this query\n",
    "1 = relevant, may be partly helpful but might contain other irrelevant content\n",
    "0 = not relevant, should never be shown for this query\n",
    "\n",
    "Assume that you are writing a report on the subject of the topic. If you would use any of the information contained in the web page in such a report, mark it 1. If the web page is primarily about the topic, or contains vital information about the topic, mark it 2. Otherwise, mark it 0.\n",
    "\n",
    "Query\n",
    "A person has typed [{query}] into a search engine.\n",
    "\n",
    "Result\n",
    "Consider the following web page.\n",
    "—BEGIN WEB PAGE CONTENT—\n",
    "{evidence}\n",
    "—END WEB PAGE CONTENT—\n",
    "\n",
    "Instructions\n",
    "Split this problem into steps:\n",
    "\n",
    "Consider the underlying intent of the search.\n",
    "Measure how well the content matches a likely intent of the query (M).\n",
    "Measure how trustworthy the web page is (T).\n",
    "Consider the aspects above and the relative importance of each, and decide on a final score (O).\n",
    "\n",
    "Produce a JSON array of scores without providing any reasoning. Example: {{\"M\": 2, \"T\": 1, \"O\": 1}}\n",
    "\n",
    "Results [{{\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70e71e11-de32-4e68-a369-7bed12d42c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "\n",
    "def extract_json(text_response):\n",
    "    # This pattern matches a string that starts with '{' and ends with '}'\n",
    "    pattern = r'\\{[^{}]*\\}'\n",
    "    matches = re.finditer(pattern, text_response)\n",
    "    json_objects = []\n",
    "    for match in matches:\n",
    "        json_str = match.group(0)\n",
    "        try:\n",
    "            # Validate if the extracted string is valid JSON\n",
    "            json_obj = json.loads(json_str)\n",
    "            json_objects.append(json_obj)\n",
    "        except json.JSONDecodeError:\n",
    "            # Extend the search for nested structures\n",
    "            extended_json_str = _json_extend_search(text_response, match.span())\n",
    "            try:\n",
    "                json_obj = json.loads(extended_json_str)\n",
    "                json_objects.append(json_obj)\n",
    "            except json.JSONDecodeError:\n",
    "                # Handle cases where the extraction is not valid JSON\n",
    "                continue\n",
    "    if json_objects:\n",
    "        return json_objects\n",
    "    else:\n",
    "        return None  # Or handle this case as you prefer\n",
    "\n",
    "def _json_extend_search(text, span):\n",
    "    # Extend the search to try to capture nested structures\n",
    "    start, end = span\n",
    "    nest_count = 0\n",
    "    for i in range(start, len(text)):\n",
    "        if text[i] == '{':\n",
    "            nest_count += 1\n",
    "        elif text[i] == '}':\n",
    "            nest_count -= 1\n",
    "            if nest_count == 0:\n",
    "                return text[start:i+1]\n",
    "    return text[start:end]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4733a72-dd63-407f-9cb3-da1ce86b8146",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "def exponential_backoff(attempt, base_delay=3, max_delay=60.0):\n",
    "    delay = base_delay * (2 ** attempt)  # exponential growth\n",
    "    delay = min(delay, max_delay)\n",
    "    jitter = random.uniform(0, delay * 0.15)  # add up to 15% jitter\n",
    "    return delay + jitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "718fbdca-d8a2-4bfd-a2b6-5822fecf195f",
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_qe_pairs:\n",
    "    import pickle\n",
    "    import os\n",
    "    import time\n",
    "    import datetime\n",
    "    import vertexai\n",
    "    from vertexai.generative_models import GenerativeModel\n",
    "    from tqdm import tqdm\n",
    "    from google import genai\n",
    "    \n",
    "    # file paths for storing results\n",
    "    root_path = \"/Users/jjr/output/rarr-rep/gemini_labels/query-evidence/\"\n",
    "    query_evidence_labels_dict_path = os.path.join(root_path, \"query_evidence_labels_dict.pkl\")\n",
    "    failed_api_calls_dict_path = os.path.join(root_path, \"failed_api_calls_dict.pkl\")\n",
    "\n",
    "    # load empty results dictionary\n",
    "    with open(\"/Users/jjr/output/rarr-rep/gemini_labels/query-evidence/query_evidence_labels_dict_empty.pkl\", \"rb\") as f:\n",
    "        query_evidence_labels_dict_empty = pickle.load(f)\n",
    "    \n",
    "    # load main results dictionary\n",
    "    with open(\"/Users/jjr/output/rarr-rep/gemini_labels/query-evidence/query_evidence_labels_dict.pkl\", \"rb\") as f:\n",
    "        query_evidence_labels_dict = pickle.load(f)\n",
    "\n",
    "    # update main results with any new q,e pairs\n",
    "    \n",
    "    # get all q,e pairs from blank\n",
    "    qe_pairs_empty = set(query_evidence_labels_dict_empty.keys())\n",
    "    # get all q,e pairs from labeled\n",
    "    qe_pairs_labeled = set(query_evidence_labels_dict.keys())\n",
    "    # add new q,e pairs to labeled\n",
    "    qe_add = qe_pairs_empty - qe_pairs_labeled\n",
    "    print(f\"adding {len(qe_add)} new q,e pairs\")\n",
    "    for (q,e) in list(qe_add):\n",
    "        query_evidence_labels_dict[q,e] = None\n",
    "    \n",
    "    # load evidence data file\n",
    "    with open(\"/Users/jjr/output/rarr-rep/gemini_labels/query-evidence/evidence_dict.pkl\", \"rb\") as f:\n",
    "        evidence_dict = pickle.load(f)\n",
    "        \n",
    "    # attempt to load any previously failed API calls; if not, start fresh.\n",
    "    if os.path.exists(failed_api_calls_dict_path):\n",
    "        with open(failed_api_calls_dict_path, \"rb\") as f:\n",
    "            failed_api_calls_dict = pickle.load(f)\n",
    "    else:\n",
    "        failed_api_calls_dict = {}\n",
    "    \n",
    "    # create list of (query, chunked_id) tuples that have None as value\n",
    "    no_eval_list = [\n",
    "        (query, chunk_id) \n",
    "        for (query, chunk_id), val in query_evidence_labels_dict.items()\n",
    "        if val is None and (query, chunk_id) not in failed_api_calls_dict\n",
    "    ]\n",
    "    \n",
    "    \n",
    "    # instantiate model\n",
    "    api_key = \"\"\n",
    "    client = genai.Client(api_key=api_key)\n",
    "    \n",
    "    max_retries = 5\n",
    "    \n",
    "    # iterate over tuples in no_eval_list\n",
    "    i = 0\n",
    "    for (query, chunk_id) in tqdm(no_eval_list, desc=\"Processing pairs\"):\n",
    "    \n",
    "        # extract evidence text using chunk_id\n",
    "        evidence_text = evidence_dict[chunk_id]\n",
    "    \n",
    "        # Populate the prompt with the required data\n",
    "        prompt_data = { \"query\": query, \"evidence\": evidence_text}\n",
    "    \n",
    "        # create llm messages\n",
    "        messages = prompt.format(**prompt_data)\n",
    "        \n",
    "        # send to gemini\n",
    "        success = False\n",
    "        last_exception = None\n",
    "        for attempt in range(max_retries):\n",
    "            try:\n",
    "                response = client.models.generate_content(model=\"gemini-1.5-pro-002\", contents=messages)\n",
    "                response_text = response.text\n",
    "                success = True\n",
    "                time.sleep(0.5)\n",
    "                break\n",
    "            except Exception as e:\n",
    "                last_exception = e\n",
    "                if hasattr(e, 'response') and getattr(e.response, 'status_code', None) == 429:\n",
    "                    delay = exponential_backoff(attempt)\n",
    "                    print(f\"HTTP 429 received. Waiting {delay:.2f} seconds before retrying.\")\n",
    "                else:\n",
    "                    delay = exponential_backoff(attempt)\n",
    "                    print(f\"API call error: {e}. Retrying in {delay:.2f} seconds.\")\n",
    "                time.sleep(delay)\n",
    "    \n",
    "        # if the API call was not successful, flag and continue (do not update the main dictionary)\n",
    "        if not success:\n",
    "            error_msg = f\"API call failed after {max_retries} attempts: {last_exception}\"\n",
    "            failed_api_calls_dict[(query, chunk_id)] = error_msg\n",
    "            continue\n",
    "        \n",
    "        try:\n",
    "            json_results = extract_json(response_text)\n",
    "            if not json_results:\n",
    "                raise ValueError(\"No valid JSON found in response\")\n",
    "            response_json = json_results[0]\n",
    "        except Exception as e:\n",
    "            print(f\"JSON parsing error for pair ({query}, {chunk_id}): {e}\")\n",
    "            failed_api_calls_dict[(query, chunk_id)] = f\"JSON parsing error: {e}\"\n",
    "            continue  \n",
    "        \n",
    "        # save label in dictionary\n",
    "        query_evidence_labels_dict[query, chunk_id] = response_json\n",
    "        i += 1\n",
    "    \n",
    "        # save progress every 1000 iterations to avoid losing good API calls\n",
    "        if i % 250 == 0:\n",
    "            print(f\"saving iteration {i}\")\n",
    "            # save timestamped version\n",
    "            timestamp = datetime.datetime.now().strftime('%d_%m_%y_%H%M')\n",
    "            with open(os.path.join(root_path, f\"query_evidence_labels_dict_{timestamp}.pkl\"), \"wb\") as f:\n",
    "                pickle.dump(query_evidence_labels_dict, f)\n",
    "            with open(os.path.join(root_path, f\"failed_api_calls_dict_{timestamp}.pkl\"), \"wb\") as f:\n",
    "                pickle.dump(failed_api_calls_dict, f)\n",
    "            # update the fixed file names so future runs will load the latest errors\n",
    "            with open(query_evidence_labels_dict_path, \"wb\") as f:\n",
    "                pickle.dump(query_evidence_labels_dict, f)\n",
    "            with open(failed_api_calls_dict_path, \"wb\") as f:\n",
    "                pickle.dump(failed_api_calls_dict, f) \n",
    "    \n",
    "    # Save any remaining work after the loop completes to the fixed file paths\n",
    "    with open(query_evidence_labels_dict_path, \"wb\") as f:\n",
    "        pickle.dump(query_evidence_labels_dict, f)\n",
    "    with open(failed_api_calls_dict_path, \"wb\") as f:\n",
    "        pickle.dump(failed_api_calls_dict, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gcloud",
   "language": "python",
   "name": "glcoud"
  },
  "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}