examples/python_inference.py

from __future__ import annotations

import argparse
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from reframr.cli import compose_generation_context
from reframr.model import ReframrModel


def main() -> None:
    parser = argparse.ArgumentParser(description="Run Reframr-RFM-v2-Base locally.")
    parser.add_argument("--model", default=str(REPO_ROOT / "model.safetensors"))
    parser.add_argument("--prompt", default="Who are you, and what makes Reframr different?")
    parser.add_argument("--system", default="")
    parser.add_argument("--max-tokens", type=int, default=120)
    parser.add_argument("--temperature", type=float, default=0.58)
    parser.add_argument("--top-k", type=int, default=64)
    parser.add_argument("--top-p", type=float, default=0.92)
    parser.add_argument("--repetition-penalty", type=float, default=1.25)
    args = parser.parse_args()

    context = compose_generation_context(args.prompt, system=args.system)
    model = ReframrModel.load(args.model)
    print(
        model.generate_text(
            context,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
            top_p=args.top_p,
            repetition_penalty=args.repetition_penalty,
        )
    )


if __name__ == "__main__":
    main()