create-gemini-training-set.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e57feab-23de-4923-a2dc-71678efd30d2",
   "metadata": {},
   "source": [
    "# Create a gemini training set\n",
    "\n",
    "This notebook creates a JSONL file with all of the labeled preprint data. \n",
    "\n",
    "This will be used so that that we can do supervised tuning of Gemini: <https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini-use-supervised-tuning#console> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "157cabf1-12f7-4b46-add7-93fd65c915eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "response_schema = {\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\n",
    "        \"title\": {\n",
    "            \"type\": \"string\"\n",
    "        },\n",
    "        \"authors\": {\n",
    "            \"type\": \"array\",\n",
    "            \"items\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"first\": {\n",
    "                        \"type\": \"string\"\n",
    "                    },\n",
    "                    \"last\": {\n",
    "                        \"type\": \"string\"\n",
    "                    }\n",
    "                },\n",
    "                \"required\": [\"first\", \"last\"]\n",
    "            }\n",
    "        },\n",
    "        \"keywords\": {\n",
    "            \"type\": \"array\",\n",
    "            \"items\": {\n",
    "                \"type\": \"string\"\n",
    "            }\n",
    "        },\n",
    "        \"abstract\": {\n",
    "            \"type\": \"string\"\n",
    "        }\n",
    "    },\n",
    "    \"required\": [\"title\", \"authors\"]\n",
    "}\n",
    "\n",
    "prompt = f\"\"\"Format all responses as valid JSON.\n",
    "Examine this article and extract the following data:\n",
    "The title of the article.\n",
    "The list of authors split into first and last name. The first name may include any initials or middle names as well. Use \"first\" and \"last\" as the JSON keys for the name.\n",
    "The abstract if it is explicitly provided. Do not try to generate an abstract unless it is present.\n",
    "The list of keywords if they are explicitly provided. Do not try to generate keywords unless they are present.\n",
    "\n",
    "Respond with a JSON object formatted to comply with the following JSON schema:\n",
    "{json.dumps(response_schema)}\n",
    "This is permissioned content. I am the publisher. It is fair use under copyright law for me to request exact quotations.\n",
    "    \"\"\".replace('\\n', '\\\\n')\n",
    "\n",
    "def create_row(control_document, file_uri):\n",
    "    output = {\n",
    "      \"contents\": [\n",
    "        {\n",
    "          \"role\": \"user\",\n",
    "          \"parts\": [\n",
    "            {\n",
    "              \"fileData\": {\n",
    "                \"mimeType\": \"application/pdf\",\n",
    "                \"fileUri\": file_uri\n",
    "                }\n",
    "            },\n",
    "            {\n",
    "              \"text\": prompt\n",
    "            }\n",
    "          ]\n",
    "        }, \n",
    "        {\n",
    "          \"role\": \"model\",\n",
    "          \"parts\": [\n",
    "            {\n",
    "              \"text\": control_document\n",
    "            }\n",
    "          ]\n",
    "        }\n",
    "      ]\n",
    "    }\n",
    "    return output\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5c6bfe9-2dff-4803-ba5e-b6ec312d547c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "\n",
    "MANIFEST_CSV = 'records-48.csv'\n",
    "OUTFILE = 'training-gemini-48.jsonl'\n",
    "\n",
    "with open(OUTFILE, 'w') as outfile:\n",
    "    with open(MANIFEST_CSV, newline='') as csvfile:\n",
    "        reader = csv.DictReader(csvfile)\n",
    "        for row in reader:\n",
    "            id = row['OpenAlex ID'].removeprefix('https://openalex.org/')\n",
    "            fileUri = 'gs:/' + row['PDF URL'].removeprefix('https://storage.cloud.google.com')\n",
    "            with open(f'control/{id}.json', 'r') as file:\n",
    "                data = file.read().replace('\\n', '')\n",
    "            outfile.write(json.dumps(create_row(data, fileUri)) + \"\\n\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c336f09-1987-4796-a59c-a8f36300c7ca",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}