MMLLM Experiments with ollama.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e675e841-ee0b-45a1-be69-e797cc5469b2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-13T10:58:07.920418Z",
     "iopub.status.busy": "2024-12-13T10:58:07.920250Z",
     "iopub.status.idle": "2024-12-13T10:58:08.245951Z",
     "shell.execute_reply": "2024-12-13T10:58:08.245606Z",
     "shell.execute_reply.started": "2024-12-13T10:58:07.920403Z"
    }
   },
   "source": [
    "Processing pipeline adapted from Taylor Arnold and Lauren Tilton's \"Explainable Search and Discovery of Visual Cultural Heritage Collections with Multimodal Large Language Models\" https://2024.computational-humanities-research.org/papers/paper28/ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9addec5f-cda4-4c91-98a7-ee94a068d4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tropy\n",
    "from dotenv import load_dotenv\n",
    "import numpy as np\n",
    "import faiss\n",
    "import base64\n",
    "from PIL import Image\n",
    "from typing import List, Tuple\n",
    "from ollama import chat\n",
    "from ollama import embeddings\n",
    "from ollama import ChatResponse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b27076b-e71c-4e1b-b7b6-073c47a8db87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration -- not currently used\n",
    "# Set ollama model\n",
    "# model='llama3.2-vision'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9113da5a-7fa5-438f-b785-ca61c5069211",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a vector database (using FAISS for simplicity)\n",
    "DIMENSION = 1024  # Embedding dimension (ensure this matches your embedding model) 3072 for text-embedding-3-large; 1024 for mxbai-embed-large\n",
    "index = faiss.IndexFlatL2(DIMENSION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7387d256-417a-42ee-b1e7-dea8dfe8cb34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping of image embeddings to Tropy metadata\n",
    "metadata_store = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e165f7eb-40ae-4879-b66b-4aa9adf36f79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to process an image and generate a caption with ollama\n",
    "def generate_caption(image_path: str) -> str:\n",
    "    \"\"\"\n",
    "    Generate a caption for an image using ollama.\n",
    "    \"\"\"\n",
    "    \n",
    "    prompt = \"Please provide a detailed plain text description of what this photograph portrays. Also provide a description of its composition as a photograph.\"\n",
    "\n",
    "    \n",
    "    response = chat(\n",
    "        model='llama3.2-vision',\n",
    "        messages=[{\n",
    "        'role': 'user',\n",
    "        'content': prompt,\n",
    "        'images': [image_path]\n",
    "        }]\n",
    "    )\n",
    "    return response['message']['content']\n",
    "\n",
    "    \n",
    "    # return response.choices[0].message.content.strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a170de-4f57-4cf9-b512-d20755c96b55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to generate embeddings for a text with ollama\n",
    "def generate_embedding(text: str) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Generate a text embedding using ollama\n",
    "    \"\"\"\n",
    "\n",
    "    response = embeddings(\n",
    "        model='mxbai-embed-large',\n",
    "        prompt=text,\n",
    "    )\n",
    "\n",
    "    return np.array(response[\"embedding\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "345d5514-b16a-44cc-a445-7bf158f75d5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve all items from Tropy\n",
    "items = tropy.get_all_items()\n",
    "print(items)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8178bbc8-36e7-47af-aecf-634c00e53e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_images():\n",
    "    \"\"\"\n",
    "    Process images, generate captions and embeddings, and store them in the vector database.\n",
    "    \"\"\"\n",
    "    global metadata_store\n",
    "\n",
    "    for item in items:\n",
    "        try:\n",
    "            photos = tropy.get_item_photos(item)\n",
    "            if not photos:\n",
    "                print(f\"No photos found for item {item}\")\n",
    "                continue\n",
    "            \n",
    "            # Ensure photos is a list or iterable\n",
    "            if not isinstance(photos, list):\n",
    "                print(f\"Unexpected format for photos in item {item}: {photos}\")\n",
    "                continue\n",
    "\n",
    "            # Get the first photo -- TODO: handle multi-image / multi-page items\n",
    "            image_path = photos[0]\n",
    "\n",
    "            if not os.path.isfile(image_path):\n",
    "                print(f\"File not found: {image_path}\")\n",
    "                continue\n",
    "\n",
    "            # Generate a caption\n",
    "            caption = generate_caption(image_path)\n",
    "            caption = caption.replace(\"\\n\", \" \") \n",
    "                        \n",
    "            # Generate an embedding\n",
    "            embedding = generate_embedding(caption)\n",
    "            \n",
    "            # Add the embedding to the vector database\n",
    "            embedding_array = np.array([embedding], dtype=np.float32)\n",
    "\n",
    "            # Validate embedding before adding\n",
    "            if embedding_array.shape[1] != DIMENSION:\n",
    "                raise ValueError(f\"Embedding dimension mismatch. Expected {DIMENSION}, got {embedding_array.shape[1]}\")\n",
    "\n",
    "            index.add(embedding_array)\n",
    "                \n",
    "            # Store metadata\n",
    "            metadata_store[len(metadata_store)] = {\n",
    "                \"file_path\": image_path, # we probably don't need this\n",
    "                \"caption\": caption,\n",
    "                \"item_id\": item  # Adding Tropy item ID for traceability back to the application\n",
    "            }\n",
    "\n",
    "            \n",
    "#            print(f\"Processed {image_path}: {caption}\")\n",
    "        except ValueError as ve:\n",
    "            print(f\"ValueError while processing item {item}: {ve}\")\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing item {item}: {e}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c2b9037-2538-4030-a149-5604d80ba91d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to search for images\n",
    "def search_images(query: str, top_k: int = 5) -> List[Tuple[str, str, str]]:\n",
    "    \"\"\"\n",
    "    Search for images using a natural language query.\n",
    "    Returns a list of tuples containing file_path, caption, and item_id.\n",
    "    \"\"\"\n",
    "    query_embedding = generate_embedding(query)\n",
    "\n",
    "    # Perform a similarity search\n",
    "    distances, indices = index.search(np.array([query_embedding], dtype=np.float32), top_k)\n",
    "\n",
    "    results = []\n",
    "    for i in range(len(indices[0])):\n",
    "        idx = indices[0][i]\n",
    "        if idx in metadata_store:\n",
    "            results.append((\n",
    "                metadata_store[idx][\"file_path\"],\n",
    "                metadata_store[idx][\"caption\"],\n",
    "                metadata_store[idx][\"item_id\"]\n",
    "            ))\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd90a464-fd2f-44d3-a1f0-6d77257153ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's go!\n",
    "process_images()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fad401a-885c-4dcc-a1b5-352401b8f436",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform a search\n",
    "query = \"Aerospace\"\n",
    "results = search_images(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e817b13-656b-4d10-a573-42768c1b3e6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display results\n",
    "for file_path, caption, item_id in results:\n",
    "    print(f\"Image: {file_path}, Caption: {caption}, Item ID: {item_id}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8c6845e-84b2-4c39-8499-c80b31b89bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "matching_item_ids = [result[2] for result in results]  # Assuming item_id is the 3rd element in the tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4016aa8-0e47-45f6-bc9b-caf7511979ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tag matching items with the search query\n",
    "for item in matching_item_ids:\n",
    "    tropy.tag_item_by_tag_name(item, [query])"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}