{ "cells": [ { "cell_type": "markdown", "id": "cell-0", "metadata": {}, "source": [ "# Create an inference function\n", "\n", "Run `data.ipynb` and `train.ipynb` first in the same local/full mode.\n", "This notebook loads their local model and tokenizer; no S3 credentials,\n", "MongoDB server, or model downloads are needed for the default demo.\n", "The tiny demo model's scores have no scientific interpretation." ] }, { "cell_type": "markdown", "id": "85c2ee1d", "metadata": {}, "source": [ "## Local setup\n", "\n", "Use Python 3.12 or 3.13 and install this checkout with `uv pip install --editable '.[notebooks]'`.\n", "See **Run notebooks locally** in the repository README for the kernel, execution order,\n", "and full-data options. Run related notebooks from the same working directory\n", "to share the default `.cache/` model store and `tracing_database.json` database.\n" ] }, { "cell_type": "markdown", "id": "cell-4", "metadata": {}, "source": [ "For a pleasant developer experience, we create some typed models that will show up in the automatically generated OpenAPI schema specification and will also provide runtime type validation." ] }, { "cell_type": "code", "execution_count": 11, "id": "cell-5", "metadata": {}, "outputs": [], "source": [ "from typing import List\n", "from pydantic import BaseModel\n", "\n", "\n", "class Attention(BaseModel):\n", " weight: float\n", " token: str\n", "\n", "\n", "class EvaluatedSentence(BaseModel):\n", " score: float\n", " text: str\n", " explanation: List[Attention]" ] }, { "cell_type": "markdown", "id": "cell-6", "metadata": {}, "source": [ "Even though `@use_model` caches the remote files locally and it also handles deserialising objects, we only use it to store a directory. In this case, it gives back a path, the path to that directory. So, we need to load the files from that folder ourselves. In order to only load it once per process, we create a small model loader helper function.\n", "\n", "> This is usually not needed, however, when we can outsmart `dill` so for optimisation purposes, we do it." ] }, { "cell_type": "code", "execution_count": 12, "id": "cell-7", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[38;5;39mFetching cached versions of scibert-highlights\u001b[0m\n", "\u001b[38;5;39mLatest version of scibert-highlights is 0 (from versions: 0)\u001b[0m\n", "\u001b[38;5;39mFile scibert-highlights-0 found in cache\u001b[0m\n" ] } ], "source": [ "from functools import lru_cache\n", "from pathlib import Path\n", "from great_ai import use_model\n", "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n", "\n", "\n", "@lru_cache(maxsize=1)\n", "def load_tokenizer_and_model(model_path):\n", " tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)\n", " model = AutoModelForSequenceClassification.from_pretrained(\n", " model_path, local_files_only=True, attn_implementation=\"eager\"\n", " )\n", " model.eval()\n", " return tokenizer, model\n", "\n", "\n", "@use_model(\"scibert-highlights\", version=\"latest\", model_kwarg_name=\"model_path\")\n", "def get_tokenizer_and_model(model_path: Path):\n", " return load_tokenizer_and_model(model_path)" ] }, { "cell_type": "markdown", "id": "cell-8", "metadata": {}, "source": [ "Finally, implement the inference function." ] }, { "cell_type": "code", "execution_count": 13, "id": "cell-9", "metadata": {}, "outputs": [], "source": [ "from great_ai import GreatAI\n", "from great_ai.utilities import clean\n", "import torch\n", "\n", "\n", "@GreatAI.create\n", "def find_highlights(sentence: str) -> EvaluatedSentence:\n", " \"\"\"Score a sentence and expose the final layer's CLS attention per token.\n", "\n", " Attention weights are a model diagnostic, not a causal explanation.\n", " \"\"\"\n", " tokenizer, model = get_tokenizer_and_model()\n", " sentence = clean(sentence, convert_to_ascii=True, remove_brackets=True)\n", " tensors = tokenizer(sentence, return_tensors=\"pt\", truncation=True, max_length=512)\n", " with torch.inference_mode():\n", " result = model(**tensors, output_attentions=True)\n", " score = torch.softmax(result.logits, dim=1)[0, 1].item()\n", " weights = result.attentions[-1][0, :, 0, :].sum(dim=0).cpu().tolist()\n", " tokens = tokenizer.convert_ids_to_tokens(tensors[\"input_ids\"][0].tolist())\n", " explanation = [\n", " Attention(token=token, weight=round(weight, 4))\n", " for token, weight in zip(tokens[1:-1], weights[1:-1])\n", " ]\n", " return EvaluatedSentence(score=score, text=sentence, explanation=explanation)" ] }, { "cell_type": "markdown", "id": "cell-10", "metadata": {}, "source": [ "A simple test to see everything works. Note that the models list is filled by the `@use_model` call even though it's not on the main inference function." ] }, { "cell_type": "code", "execution_count": 14, "id": "cell-11", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "W1028 12:49:11.667000 52910 torch/distributed/elastic/multiprocessing/redirects.py:35] NOTE: Redirects are currently not supported in MacOs.\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "6ef4b00220304248a94035042cb5fa15", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Loading weights: 0%| | 0/25 [00:00