File size: 1,433 Bytes
1b88863 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """Native JEV JSON-lines inference; accepts state/question or state/questions."""
import argparse
import json
from pathlib import Path
import sys
ROOT=Path(__file__).resolve().parent
sys.path.insert(0,str(ROOT/'runtime'))
from maincode_jev_serve.model import DecisionModel, answer
from maincode_jev_serve.decide import decide
def main():
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--checkpoint',default=str(ROOT))
p.add_argument('--device',default='cuda:0')
p.add_argument('--max-tokens',type=int,default=131072)
p.add_argument('--token-budget',type=int,default=131072)
p.add_argument('--batch-size',type=int,default=64)
args=p.parse_args()
model=DecisionModel(checkpoint=args.checkpoint,device=args.device)
for line in sys.stdin:
if not line.strip():continue
row=json.loads(line)
questions=row.get('questions')
if questions is None:questions={'answer':row['question']}
probs,tokens=decide(model,row['state'],questions,temperature=model.temperature,
max_tokens=args.max_tokens,token_budget=args.token_budget,batch_size=args.batch_size,
images=row.get('images',()))
print(json.dumps({'model':'MATILDA JEV by Maincode','answers':{k:answer(questions[k],v) for k,v in probs.items()},
'usage':{'input_tokens':tokens}},ensure_ascii=False),flush=True)
if __name__=='__main__':main()
|