import json, sys, time, urllib.request
from concurrent.futures import ThreadPoolExecutor
cfgs = {
 "topp_temp": {"top_k":0,"top_p":0.95,"min_p":0,"samplers":["top_p","temperature"]},
 "topk50_topp": {"top_k":50,"top_p":0.95,"min_p":0,"samplers":["top_k","top_p","temperature"]},
 "temp_only": {"top_k":0,"top_p":1.0,"min_p":0,"samplers":["temperature"]},
 "greedy": {"temperature":0},
}
def one(cfg, i):
    b={"prompt":"Write a long story about a dragon number %d."%i,"n_predict":384,"ignore_eos":True,"temperature":0.6,"seed":i,"repeat_penalty":1.0,**cfg}
    r=json.loads(urllib.request.urlopen(urllib.request.Request("http://localhost:8080/completion",data=json.dumps(b).encode(),headers={"Content-Type":"application/json"})).read())
    return r["tokens_predicted"]
for name,cfg in cfgs.items():
    t=time.time()
    with ThreadPoolExecutor(48) as ex: n=sum(ex.map(lambda i: one(cfg,i), range(48)))
    print(name, round(n/(time.time()-t)), "tok/s", flush=True)