experiments/llm-inference-api/Extract-metadata-LLM-API.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1a118ebe-d684-4b2a-a8ad-a7f1df5e3d04",
   "metadata": {},
   "source": [
    "# Extract metadata from a PDF document using a LLM API service\n",
    "\n",
    "This notebook demonstrates how to extract metadata from a PDF document using an API service that provides access to a fine-tuned LLM that performs the main extraction task.\n",
    "\n",
    "For testing, you need the following:\n",
    "\n",
    "1. Install [llama.cpp](https://github.com/ggerganov/llama.cpp) on your computer (on Linux, using CPU only: `git clone` the repository and run `make` to compile it)\n",
    "2. Download a fine-tuned model such as [NatLibFi/Qwen2-0.5B-Instruct-FinGreyLit-GGUF](https://huggingface.co/NatLibFi/Qwen2-0.5B-Instruct-FinGreyLit-GGUF) in GGML format i.e. [Qwen2-0.5B-Instruct-FinGreyLit-Q4_K_M.gguf\n",
    "](https://huggingface.co/NatLibFi/Qwen2-0.5B-Instruct-FinGreyLit-GGUF/blob/main/Qwen2-0.5B-Instruct-FinGreyLit-Q4_K_M.gguf)\n",
    "3. Start the llama.cpp server using the GGUF model: `./llama-server -m Qwen2-0.5B-Instruct-FinGreyLit-Q4_K_M.gguf` and leave it running\n",
    "4. Install the Python dependencies that this notebook needs from `requirements.txt`\n",
    "5. Now you can run this notebook! Adjust LLM_API_URL below accordingly if your server is not running on `localhost:8080`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "db9c790b-8648-445e-8163-1127e1f0e9c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set config\n",
    "\n",
    "LLM_API_URL = \"http://localhost:8080/chat/completions\"  # local llama.cpp server\n",
    "MODEL_NAME = \"xxx\"  # llama-server doesn't care about this, it will simply use the model file you gave it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "922390a9-ad8c-4cdf-9016-f2e694be0cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary libraries\n",
    "\n",
    "import requests\n",
    "import json\n",
    "import io\n",
    "import fitz\n",
    "import regex\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9aa8e390-5eca-4aba-92bf-a8f7622320da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"language\": \"eng\", \"title\": \"Machine Learning Poems\", \"creator\": [\"Yoon, Jeonghee\"], \"year\": \"2021\", \"type_coar\": \"book\"}\n"
     ]
    }
   ],
   "source": [
    "# test the language model using a simple request\n",
    "\n",
    "def llm_request(message, system_prompt=\"You are a helpful assistant.\", temperature=0.0, max_tokens=50):\n",
    "\n",
    "    headers = {\n",
    "        \"Content-Type\": \"application/json\",\n",
    "    #    \"Authorization\": f\"Bearer None\"\n",
    "    }\n",
    "\n",
    "    data = {\n",
    "        \"model\": MODEL_NAME,\n",
    "        \"messages\": [\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": message},\n",
    "        ],\n",
    "        \"temperature\": temperature,\n",
    "        \"max_tokens\": max_tokens,\n",
    "        \"response_format\": {\"type\": \"json_object\"}\n",
    "    }\n",
    "\n",
    "    response = requests.post(LLM_API_URL, headers=headers, json=data)\n",
    "\n",
    "    if response.status_code == 200:\n",
    "        #print(\"Response from OpenAI:\", response.json())\n",
    "        return response.json()['choices'][0]['message']['content']\n",
    "    else:\n",
    "        print(\"Error:\", response.status_code, response.text)\n",
    "\n",
    "print(llm_request(\"Write some JSON that describes a poem collection about machine learning\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8d16423d-59c0-44c4-b025-dfaff5485b5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Please enter a PDF URL:  https://www.doria.fi/bitstream/handle/10024/188075/Tekoa%cc%88lykahvit_%20Extracting%20metadata%20using%20LLMs.pdf?sequence=1&isAllowed=y\n"
     ]
    }
   ],
   "source": [
    "# ask for PDF URL\n",
    "\n",
    "url = input(\"Please enter a PDF URL: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b7e8ecba-e2f4-4b68-9aca-424ffbdcf16b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "text length: 1232 characters\n",
      "\n",
      "{\"pdfinfo\": {\"title\": \"Teko\\u00e4lykahvit: Extracting metadata using LLMs\"}, \"pages\": [{\"page\": 0, \"text\": \"Extracting metadata from grey literature using large language models\\nOsma Suominen\\nTeko\\u00e4lykahvit 1.11.2023\\nPerustuu SWIB23-konferenssin salamaesitykseen 12.9.2023\"}, {\"page\": 1, \"text\": \"Grey literature? reports working papers government documents white papers preprints theses \\u2026 semi-formal non-commercial\\nPDFs published on the web \\u2013 lots of them!\\nImage made with DreamStudio (based on Stable Diffusion)\\nPrompt: \\\"A big pile of papers, reports, documents, PDF \\ufb01les, word documents, powerpoint slides, posters, articles and books forming a wave in the style of Hokusai\\\"\"}, {\"page\": 5, \"text\": \"First 5 pages of text from PDF contributor/faculty: fi=School of\\nFine-tuned\\nWasaensia relation/numberinseries: 500\"}, {\"page\": 6, \"text\": \"Example of LLM extracted metadata\\nDiff view: human vs. LLM generated\"}, {\"page\": 7, \"text\": \"What we found out\"}, {\"page\": 11, \"text\": \"Results\\nIt's looking pretty good so far.\\nThe data set must still be improved.\\nNo results to publish yet, stay tuned!\"}, {\"page\": 12, \"text\": \"Thank you!\\nOsma Suominen osma.suominen@helsinki.\\ufb01 @osma@sigmoid.social\"}]}\n"
     ]
    }
   ],
   "source": [
    "# download the PDF and extract the relevant text\n",
    "\n",
    "# settings for text extraction\n",
    "PAGES = [0, 1, 2, 3, 4, 5, 6, 7, -2, -1]  # pages to analyze: first 8 pages + last page\n",
    "THRESHOLD = 100                       # paragraphs shorter than this will always be kept\n",
    "LONG_PARAGRAPH_PAGES = [0, 1]         # on first two pages, some long paragraphs are accepted\n",
    "LONG_PARAGRAPH_MAX = 2                # how many long paragraphs to keep on the first two pages\n",
    "PDF_METADATA_SKIP = {'format', 'creator', 'producer'}  # PDF metadata fields not to include in extracted text\n",
    "\n",
    "def download_and_open_pdf(url):\n",
    "    response = requests.get(url)\n",
    "    pdf_stream = io.BytesIO(response.content)\n",
    "    return fitz.open(stream=pdf_stream, filetype=\"pdf\")\n",
    "\n",
    "def extract_content(pdf):\n",
    "    \"\"\"extract and return the pdfinfo metadata and the first few pages of text (and last page) from the given PDF file\"\"\"\n",
    "\n",
    "    pdfinfo = {}\n",
    "    pages = []\n",
    "    \n",
    "    for key in pdf.metadata.keys():\n",
    "        if key not in PDF_METADATA_SKIP and pdf.metadata.get(key):\n",
    "            pdfinfo[key] = pdf.metadata.get(key)\n",
    "\n",
    "    for page in PAGES:\n",
    "        if page > len(pdf) - 2:\n",
    "            continue\n",
    "\n",
    "        texts = []\n",
    "        text = pdf[page].get_text(sort=True)\n",
    "        # Use regular expression to split text into paragraphs\n",
    "        # Delimiter: newline(s) followed by an upper case character\n",
    "        paragraphs = regex.split(r'\\n+(?=\\p{Lu})', text, flags=re.UNICODE)\n",
    "        long_paragraph_count = 0\n",
    "\n",
    "        for paragraph in paragraphs:\n",
    "            paragraph = \" \".join(paragraph.strip().split())\n",
    "\n",
    "            if '.....' in paragraph or '. . . . .' in paragraph: # looks like a ToC entry, skip it\n",
    "                continue\n",
    "            elif len(paragraph) < THRESHOLD:  # short paragraph, keep it\n",
    "                texts.append(paragraph)\n",
    "            elif page in LONG_PARAGRAPH_PAGES and long_paragraph_count < LONG_PARAGRAPH_MAX:\n",
    "                # allow some long paragraphs on the first two pages\n",
    "                long_paragraph_count += 1\n",
    "                texts.append(paragraph)\n",
    "            else:  # must be a long paragraph, skip it\n",
    "                pass\n",
    "        text = '\\n'.join(texts)\n",
    "        if text:\n",
    "            pages.append({\"page\": pdf[page].number, \"text\": text})\n",
    "\n",
    "    return {\"pdfinfo\": pdfinfo, \"pages\": pages}\n",
    "\n",
    "pdf = download_and_open_pdf(url)\n",
    "\n",
    "doc_json = json.dumps(extract_content(pdf))\n",
    "print(f\"text length: {len(doc_json)} characters\")\n",
    "print()\n",
    "print(doc_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8886a7f7-640e-4802-9e2d-cd70c9f9b580",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "{   'creator': ['Suominen, Osma'],\n",
      "    'language': 'eng',\n",
      "    'publisher': ['University of Helsinki'],\n",
      "    'title': 'Tekoälykahvit : Extracting metadata using large language models',\n",
      "    'type_coar': 'research article',\n",
      "    'year': '2023'}\n",
      "\n",
      "CPU times: user 5.22 ms, sys: 0 ns, total: 5.22 ms\n",
      "Wall time: 5.18 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# submit the text to the LLM and display results\n",
    "\n",
    "import pprint\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "\n",
    "SYSTEM_PROMPT = \"You are a skilled librarian specialized in meticulous cataloguing of digital documents.\"\n",
    "INSTRUCTION = \"Extract metadata from this document. Return as JSON.\"\n",
    "MAX_TOKENS = 1024\n",
    "\n",
    "message = f\"{INSTRUCTION}\\n\\n{doc_json}\"\n",
    "response = llm_request(message, system_prompt=SYSTEM_PROMPT, max_tokens=1024)\n",
    "#print(response)\n",
    "extracted_data = json.loads(response)\n",
    "print()\n",
    "pp.pprint(extracted_data)\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7ba3fab-598a-4c74-a4ce-c9672b94e203",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}