predict.py

# /// script
# requires-python = ">=3.12"
# dependencies = ["setfit==1.2.0", "sentence-transformers==6.0.1", "transformers==5.16.1", "torch==2.14.0", "scikit-learn==1.9.0", "huggingface-hub==1.30.0"]
# ///
"""Predict document purpose: uv run --python 3.12 predict.py document.txt."""

import argparse
import hashlib
import json
from pathlib import Path

from setfit import SetFitModel
from transformers import AutoTokenizer

MODEL = "davanstrien/finepdfs-edu-purpose-classifier"
PREPARATION_TOKENIZER = "ibm-granite/granite-embedding-small-english-r2"
TOKENIZER_REVISION = "2ab6fa8ea2d674564defd37171ae19079b864b33"

def sha(data):
    return hashlib.sha256(data).hexdigest()


def prepare(text, tokenizer):
    tokens = tokenizer.encode(text, add_special_tokens=False, truncation=False, verbose=False)
    excerpt, mode = text, "full"
    if len(tokens) > 990:
        midpoint = (len(tokens) - 330) // 2
        windows = [tokens[:330], tokens[midpoint : midpoint + 330], tokens[-330:]]
        excerpt = "\n[...]\n".join(
            tokenizer.decode(part, skip_special_tokens=True, clean_up_tokenization_spaces=False)
            for part in windows
        )
        mode = "beginning_middle_end_330_each"
    prepared = "classification: " + excerpt
    return prepared, dict(
        prepared_sha256=sha(prepared.encode()),
        original_preparation_tokens=len(tokens),
        excerpt_mode=mode,
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("text_file", type=Path)
    parser.add_argument("--device", default="cpu", choices=("cpu", "cuda"))
    args = parser.parse_args()
    tokenizer = AutoTokenizer.from_pretrained(PREPARATION_TOKENIZER, revision=TOKENIZER_REVISION)
    model = SetFitModel.from_pretrained(MODEL, device=args.device)
    model.model_body.float().eval()
    model.model_body.max_seq_length = 1024
    model.model_body.default_prompt_name = None
    prepared, metadata = prepare(args.text_file.read_text(), tokenizer)
    vectors = model.model_body.encode([prepared], normalize_embeddings=False, convert_to_numpy=True)
    scores = model.model_head.predict_proba(vectors)[0]
    labels = model.model_head.classes_.tolist()
    print(json.dumps({
        "predicted_purpose": labels[int(scores.argmax())],
        "scores": dict(zip(labels, scores.tolist())),
        **metadata,
    }, indent=2))


if __name__ == "__main__":
    main()