examples/basic.py

"""Run all three judgment types. Execute from anywhere after installing requirements."""
import argparse
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from openjudgement import OpenJudgement

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default=str(ROOT))
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--request', type=Path, default=ROOT/'examples/request.json')
    args = parser.parse_args()
    request = json.loads(args.request.read_text())
    judge = OpenJudgement.from_pretrained(args.model, device=args.device)
    result = judge.system_one(request['state'], request['questions'])
    print(json.dumps(result, indent=2, ensure_ascii=False))

if __name__ == '__main__':
    main()