File size: 1,181 Bytes
c69aaec | 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 | """Read native JEV state/question JSON lines and write native decision answers."""
import argparse,json,sys
from pathlib import Path
from jev_fp4 import FP4DecisionModel
from kev.model import answer
from kev.decide import decide
def main():
parser=argparse.ArgumentParser()
parser.add_argument('--checkpoint',default=str(Path(__file__).resolve().parent))
parser.add_argument('--device',default='cuda:0')
args=parser.parse_args()
model=FP4DecisionModel(args.checkpoint,device=args.device)
for line in sys.stdin:
if not line.strip():continue
row=json.loads(line)
if 'questions' in row:
probabilities,tokens=decide(model,row['state'],row['questions'],temperature=model.temperature,
max_tokens=131072,token_budget=131072,batch_size=64,images=row.get('images',()))
result={'answers':{k:answer(row['questions'][k],v) for k,v in probabilities.items()},'usage':{'input_tokens':tokens}}
else:
result=answer(row['question'],model.predict([row],batch_size=1)[0])
print(json.dumps(result,ensure_ascii=False),flush=True)
if __name__=='__main__':main()
|