inference.py

#!/usr/bin/env python
"""Run compliantLLM from its Hugging Face repository directory."""

import argparse
from pathlib import Path

from transformers import AutoModel, AutoTokenizer


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("prompt")
    args = parser.parse_args()

    repository = Path(__file__).resolve().parent
    tokenizer = AutoTokenizer.from_pretrained(repository, trust_remote_code=True)
    model = AutoModel.from_pretrained(repository, trust_remote_code=True).eval()
    inputs = tokenizer(
        args.prompt,
        return_tensors="pt",
        truncation=True,
        max_length=1024,
    )
    output_ids = model.generate(**inputs)[0]
    print(model.decode_output(output_ids))


if __name__ == "__main__":
    main()