File size: 2,422 Bytes
3055b70 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # /// 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()
|