fine_tune.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7bd8fa2-4490-4f96-846e-e112523397ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "import pandas as pd\n",
    "import json\n",
    "\n",
    "from predibase import PredibaseClient\n",
    "from neo4j import GraphDatabase, RoutingControl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "069d2b4c-d68d-461a-a79d-fba967f398a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "neo4j_uri = os.environ.get('NEO4J_URI')\n",
    "neo4j_username = os.environ.get('NEO4J_USERNAME')\n",
    "neo4j_password = os.environ.get('NEO4J_PASSWORD')\n",
    "neo4j_auth =  (neo4j_username, neo4j_password)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e3374e-7afb-49f5-a2e6-58903836a617",
   "metadata": {},
   "outputs": [],
   "source": [
    "open_ai_api_key = os.environ.get('OPENAI_API_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af034b1a-eb44-4975-912d-298db177929f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pc = PredibaseClient()\n",
    "driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_auth))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef8d52ab-77fd-4e09-bfe6-8ff0467351f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# helper function to execute the predibase driver\n",
    "\n",
    "def execute_query(driver, query):\n",
    "    with driver.session() as session:\n",
    "        def _execute(tx):\n",
    "            result = tx.run(query)\n",
    "            return [record.data() for record in result]\n",
    "        return session.execute_read(_execute)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a425b51-ef21-451e-bb05-7452e7f52651",
   "metadata": {},
   "source": [
    "#### Load and preview training data\n",
    "\n",
    "Using the training data created synthetically using `generate_synthetic_data`directory in this repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984d537e-0b7e-4b18-9356-bfebc1e9a4ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_data = 'training-data-twitter.csv'\n",
    "df = pd.read_csv(training_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f0a4248-bec7-4bf4-960e-e9607f5fa5d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Number of rows: {len(df)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c74b65f-f3a1-4147-b085-8ab521ac8bd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "945f9a4b-9c49-42d0-acd8-06cdfc653bfd",
   "metadata": {},
   "source": [
    "#### Upload training data to Predibase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e9ae528-2255-4979-82b6-d9c910fcd89f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = pc.upload_dataset(training_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66da14da-976b-49ed-b15e-b2f614bb15af",
   "metadata": {},
   "source": [
    "#### Configure and start the training job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ffbee44-4457-4cae-ba20-8a9874e5c75d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defines the template used to prompt the model for each example\n",
    "\n",
    "prompt_template = \"\"\"Below is an instruction that describes a task, paired with an input\n",
    "    that may provide further context. Write a response that appropriately\n",
    "    completes the request.\n",
    "\n",
    "    ### Instruction: {instruction}\n",
    "\n",
    "    ### Input: {input}\n",
    "\n",
    "    ### Response:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe492e64-f6b7-4fce-8a99-3c075e925ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specifies the Huggingface LLM you want to fine-tune\n",
    "# Kick off a fine-tuning job on the uploaded dataset\n",
    "llm = pc.LLM(\"hf://meta-llama/Llama-2-7b-hf\")\n",
    "job = llm.finetune(\n",
    "    prompt_template=prompt_template,\n",
    "    target=\"output\",\n",
    "    dataset=dataset,\n",
    ")\n",
    "\n",
    "# Wait for the job to finish and get training updates and metrics\n",
    "model = job.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c79fe236-c0f0-4567-a01b-77b305011362",
   "metadata": {},
   "source": [
    "#### Load new model with adapter\n",
    "\n",
    "Wait for the model to have finished training in the previous step before setting up the deployment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359ab838-7c89-4af9-9b1d-251740d41680",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_deployment = pc.LLM(\"pb://deployments/llama-2-7b\")\n",
    "\n",
    "# Set the model \n",
    "\n",
    "model = pc.get_model(\"Llama-2-7b-hf-training-data-twitter\")\n",
    "\n",
    "# Set the adapapter to our newly fine tuned model\n",
    "\n",
    "adapter_deployment = base_deployment.with_adapter(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce38caeb-7138-433c-8312-c428bed04c72",
   "metadata": {},
   "source": [
    "#### Test the new model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1802618-63f1-4e91-91ff-1b79adbfe10a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# introspect schema from neo4j\n",
    "query = \"CALL apoc.meta.graph()\"\n",
    "meta_graph = execute_query(driver, query)\n",
    "meta_graph_str = json.dumps(meta_graph)\n",
    "\n",
    "result = adapter_deployment.prompt(\n",
    "    {\n",
    "      \"instruction\": f\"Given this schema, write a Cypher query that returns the data I am looking for. Schema:  {meta_graph_str}\",\n",
    "      \"input\": \"I am neo4j. Find the hashtags used in my tweets that have the most favourites.\"\n",
    "    },\n",
    "    max_new_tokens=256)\n",
    "print(result.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72c36839-bd97-4890-b4b6-01dc0896f99c",
   "metadata": {},
   "source": [
    "#### Deploy new model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e07d14-4f61-41ed-9108-06c53e057e05",
   "metadata": {},
   "outputs": [],
   "source": [
    "finetuned_llm = model.deploy(\"training-data-twitter\").get()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bfe371c",
   "metadata": {},
   "source": [
    "If you have already fine tuned the model and want to use it for prompts: \n",
    "- Comment out the above \n",
    "- Uncomment the below "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ec6e17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the adapter to use, which is the model you have already fine-tuned.\n",
    "# model = pc.get_model(\"Llama-2-7b-hf-training-data-twitter\")\n",
    "# adapter_deployment = base_deployment.with_adapter(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80d179a6-7477-426b-9a9c-1cfd7289da36",
   "metadata": {},
   "source": [
    "#### Test newly deployed model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "430d8efe-181d-4b31-860c-6bd52c2c96af",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = finetuned_llm.prompt(\n",
    "    {\n",
    "        \"instruction\": f\"Given this schema, write a Cypher query that returns the data I am looking for.  Schema:  {meta_graph_str}\",\n",
    "      \"input\": \"I am neo4j. How many of my tweets did 'nsmith_piano' reply to?\"\n",
    "    },\n",
    "    max_new_tokens=256)\n",
    "\n",
    "\n",
    "# Uncomment to use adapter\n",
    "\n",
    "# result = adapter_deployment.prompt(\n",
    "#     {\n",
    "#       \"instruction\": f\"Given this schema, write a Cypher query that returns the data I am looking for. Schema:  {meta_graph_str}\",\n",
    "#       \"input\": \"I am neo4j. Find the hashtags used in my tweets that have the most favourites.\"\n",
    "#     },\n",
    "#     max_new_tokens=256)\n",
    "\n",
    "print(result.response)\n",
    "final_query = result.response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f3ebd7-fc22-4813-9bff-105d658cf521",
   "metadata": {},
   "source": [
    "#### Call Neo4J with the query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf15e5c5-0069-477c-9a7c-b2763e410ac2",
   "metadata": {},
   "outputs": [],
   "source": [
    "neo4j_result = execute_query(driver, final_query)\n",
    "\n",
    "print(neo4j_result)"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}