example.py

"""Minimal end-to-end example for Quicksort-fr/T-CRED-SL."""

from __future__ import annotations

import json

import torch
from transformers import AutoModel, AutoTokenizer


MODEL_ID = "Quicksort-fr/T-CRED-SL"


def main() -> None:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
    model.to("cuda" if torch.cuda.is_available() else "cpu")

    records = [
        {
            "input_id": "answer-example",
            "task": "answer",
            "question": "Who led the organization in 2020?",
            "query_time_or_interval": "2020",
            "temporal_operator": "during",
            "reference_answers": ["Alex Morgan"],
            "candidate_or_claim": "Alex Morgan",
        },
        {
            "input_id": "support-example",
            "task": "support",
            "question": "Who led the organization in 2020?",
            "candidate_or_claim": "Alex Morgan led the organization in 2020.",
            "evidence_passages": [
                {
                    "evidence_id": "e1",
                    "text": "Alex Morgan served as the organization's director from 2019 to 2021.",
                }
            ],
        },
    ]
    predictions = model.predict(tokenizer, records, batch_size=8)
    print(json.dumps(predictions, indent=2))


if __name__ == "__main__":
    main()