example.py

"""Offline contracted scientific decision. MycelAI, Apache-2.0."""
import argparse
import json
from pathlib import Path

import torch
from transformers import AutoModel, AutoTokenizer

parser = argparse.ArgumentParser()
parser.add_argument("--model", type=Path, default=Path(__file__).resolve().parent)
parser.add_argument("--device", default="cuda")
args = parser.parse_args()
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True, local_files_only=True)
model = AutoModel.from_pretrained(args.model, trust_remote_code=True,
                                  local_files_only=True, dtype=torch.bfloat16)
model.prepare_for_inference(args.device)
result = model.predict(tokenizer, task="study_approach", text=(
    "Title: A synthetic benchmark for sorting algorithms\n"
    "Abstract: We compare algorithmic time complexity using generated arrays "
    "and prove a bound for a new sorting method."))
print(json.dumps(result, indent=2))