inference_mmbert_hf_example.py

#!/usr/bin/env python3
"""
Concise inference: load HF bundle and predict on one text.
Usage:
  python inference_mmbert_hf_example.py [model_dir] [text]
Defaults:
  model_dir = .
  text      = simple Dutch example
"""

import sys, torch
from model import DualHeadModel


def main():
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "."
    text = sys.argv[2] if len(sys.argv) > 2 else (
        "Het is echt NIET te doen hier!!! Door dat hele filmfestival zijn er elke avond mensen aan het schreeuwen en harde muziek tot laat Ik kan gewoon niet meer slapen  Hoe is dit ooit goedgekeurd zo vlak na de feestdagen?????? Heb al beelden gemaakt als bewijs, kan ik die ergens heen sturen?? Het moet toch snellre opgelost kunnen worden dan dit, het duurt allemaal veel te lang Kunnen jullie dr ff naar kijken????"
    )

    device = torch.device("cuda" if torch.cuda.is_available() else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu"))

    # Load model
    model, tokenizer, config = DualHeadModel.from_pretrained(model_dir, device=device)

    # Encode and predict
    enc = tokenizer(text, truncation=True, padding="max_length", max_length=config["max_length"], return_tensors="pt")
    input_ids, attn = enc["input_ids"].to(device), enc["attention_mask"].to(device)
    on_probs, be_probs = model.predict(input_ids, attn)
    on_probs, be_probs = on_probs[0].cpu(), be_probs[0].cpu()

    # Top-5 per head (probability)
    def topk(probs, names, k=5):
        idx = torch.topk(probs, k=min(k, len(probs))).indices.tolist()
        return [(names[i], float(probs[i])) for i in idx]

    print(f"Onderwerp top-5: {[f'{n}: {p:.3f}' for n,p in topk(on_probs, config['labels']['onderwerp'])]}")
    print(f"Beleving top-5:  {[f'{n}: {p:.3f}' for n,p in topk(be_probs, config['labels']['beleving'])]}")
    print(f"Device: {device} | max_length: {config['max_length']} | model_dir: {model_dir}")


if __name__ == "__main__":
    main()