Create an inference function¶

Run data.ipynb and train.ipynb first in the same local/full mode. This notebook loads their local model and tokenizer; no S3 credentials, MongoDB server, or model downloads are needed for the default demo. The tiny demo model's scores have no scientific interpretation.

Local setup¶

Use Python 3.12 or 3.13 and install this checkout with uv pip install --editable '.[notebooks]'. See Run notebooks locally in the repository README for the kernel, execution order, and full-data options. Run related notebooks from the same working directory to share the default .cache/ model store and tracing_database.json database.

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.

In [1]:

from typing import List
from pydantic import BaseModel


class Attention(BaseModel):
    weight: float
    token: str


class EvaluatedSentence(BaseModel):
    score: float
    text: str
    explanation: List[Attention]

from typing import List from pydantic import BaseModel class Attention(BaseModel): weight: float token: str class EvaluatedSentence(BaseModel): score: float text: str explanation: List[Attention]

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.

This is usually not needed, however, when we can outsmart dill so for optimisation purposes, we do it.

In [2]:

from functools import lru_cache
from pathlib import Path
from great_ai import use_model
from transformers import AutoModelForSequenceClassification, AutoTokenizer


@lru_cache(maxsize=1)
def load_tokenizer_and_model(model_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_path, local_files_only=True, attn_implementation="eager"
    )
    model.eval()
    return tokenizer, model


@use_model("scibert-highlights", version="latest", model_kwarg_name="model_path")
def get_tokenizer_and_model(model_path: Path):
    return load_tokenizer_and_model(model_path)

from functools import lru_cache from pathlib import Path from great_ai import use_model from transformers import AutoModelForSequenceClassification, AutoTokenizer @lru_cache(maxsize=1) def load_tokenizer_and_model(model_path): tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True) model = AutoModelForSequenceClassification.from_pretrained( model_path, local_files_only=True, attn_implementation="eager" ) model.eval() return tokenizer, model @use_model("scibert-highlights", version="latest", model_kwarg_name="model_path") def get_tokenizer_and_model(model_path: Path): return load_tokenizer_and_model(model_path)

Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️

Cannot find credentials files, defaulting to using ParallelTinyDbDriver

The selected tracing database (ParallelTinyDbDriver) is not recommended for production

Cannot find credentials files, defaulting to using LargeFileLocal

GreatAI (v0.1.11): configured ✅

  🔩 tracing_database: ParallelTinyDbDriver

  🔩 large_file_implementation: LargeFileLocal

  🔩 is_production: False

  🔩 should_log_exception_stack: True

  🔩 prediction_cache_size: 512

  🔩 dashboard_table_size: 50

You still need to check whether you follow all best practices before trusting your deployment.

> Find out more at https://se-ml.github.io/practices

Fetching cached versions of scibert-highlights

Latest version of scibert-highlights is 0 (from versions: 0)

File scibert-highlights-0 found in cache

Finally, implement the inference function.

In [3]:

from great_ai import GreatAI
from great_ai.utilities import clean
import torch


@GreatAI.create
def find_highlights(sentence: str) -> EvaluatedSentence:
    """Score a sentence and expose the final layer's CLS attention per token.

    Attention weights are a model diagnostic, not a causal explanation.
    """
    tokenizer, model = get_tokenizer_and_model()
    sentence = clean(sentence, convert_to_ascii=True, remove_brackets=True)
    tensors = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=512)
    with torch.inference_mode():
        result = model(**tensors, output_attentions=True)
        score = torch.softmax(result.logits, dim=1)[0, 1].item()
        weights = result.attentions[-1][0, :, 0, :].sum(dim=0).cpu().tolist()
    tokens = tokenizer.convert_ids_to_tokens(tensors["input_ids"][0].tolist())
    explanation = [
        Attention(token=token, weight=round(weight, 4))
        for token, weight in zip(tokens[1:-1], weights[1:-1])
    ]
    return EvaluatedSentence(score=score, text=sentence, explanation=explanation)

from great_ai import GreatAI from great_ai.utilities import clean import torch @GreatAI.create def find_highlights(sentence: str) -> EvaluatedSentence: """Score a sentence and expose the final layer's CLS attention per token. Attention weights are a model diagnostic, not a causal explanation. """ tokenizer, model = get_tokenizer_and_model() sentence = clean(sentence, convert_to_ascii=True, remove_brackets=True) tensors = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=512) with torch.inference_mode(): result = model(**tensors, output_attentions=True) score = torch.softmax(result.logits, dim=1)[0, 1].item() weights = result.attentions[-1][0, :, 0, :].sum(dim=0).cpu().tolist() tokens = tokenizer.convert_ids_to_tokens(tensors["input_ids"][0].tolist()) explanation = [ Attention(token=token, weight=round(weight, 4)) for token, weight in zip(tokens[1:-1], weights[1:-1]) ] return EvaluatedSentence(score=score, text=sentence, explanation=explanation)

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.

In [4]:

if __name__ == "__main__":
    from fastapi.testclient import TestClient

    first = find_highlights("Our solution has outperformed the state-of-the-art.")
    second = find_highlights("Their solution did not perform well.")
    assert 0 <= first.output.score <= 1
    assert first.output.explanation
    print(first.output)
    print(second.output)
    with TestClient(find_highlights.app) as client:
        response = client.post(
            "/predict", json={"sentence": "Our method improves accuracy."}
        )
        response.raise_for_status()
        assert 0 <= response.json()["output"]["score"] <= 1

if __name__ == "__main__": from fastapi.testclient import TestClient first = find_highlights("Our solution has outperformed the state-of-the-art.") second = find_highlights("Their solution did not perform well.") assert 0 <= first.output.score <= 1 assert first.output.explanation print(first.output) print(second.output) with TestClient(find_highlights.app) as client: response = client.post( "/predict", json={"sentence": "Our method improves accuracy."} ) response.raise_for_status() assert 0 <= response.json()["output"]["score"] <= 1

W1028 22:48:11.714000 21400 .venv/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/redirects.py:35] NOTE: Redirects are currently not supported in MacOs.

score=0.4977932274341583 text='Our solution has outperformed the state-of-the-art.' explanation=[Attention(weight=0.1332, token='our'), Attention(weight=0.133, token='[UNK]'), Attention(weight=0.1332, token='[UNK]'), Attention(weight=0.1329, token='[UNK]'), Attention(weight=0.134, token='the'), Attention(weight=0.1342, token='[UNK]'), Attention(weight=0.1333, token='[UNK]'), Attention(weight=0.1332, token='[UNK]'), Attention(weight=0.1327, token='[UNK]'), Attention(weight=0.1316, token='the'), Attention(weight=0.1339, token='[UNK]'), Attention(weight=0.1325, token='[UNK]'), Attention(weight=0.1347, token='.')]
score=0.49779173731803894 text='Their solution did not perform well.' explanation=[Attention(weight=0.2217, token='[UNK]'), Attention(weight=0.2211, token='[UNK]'), Attention(weight=0.2215, token='[UNK]'), Attention(weight=0.221, token='[UNK]'), Attention(weight=0.2239, token='[UNK]'), Attention(weight=0.223, token='[UNK]'), Attention(weight=0.2243, token='.')]

In this case, the service is built as a docker image, pushed to our image registry and subsequent rolling update is performed in the production cluster. To check out the Dockerimage, go to the additional files page.