import json, time, urllib.request
from concurrent.futures import ThreadPoolExecutor
base = {"n_predict":384,"ignore_eos":True,"temperature":0.6,"top_k":50,"top_p":0.95,"min_p":0,"samplers":["top_k","top_p","temperature"]}
extra = {"repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0, "dry_multiplier": 0.0, "xtc_probability": 0.0, "typical_p": 1.0, "cache_prompt": False}
variants = {"base": {}, "all_extra": extra, **{k: {k: v} for k, v in extra.items()}}
def one(cfg, i):
    b = {"prompt": "Write a long story about a dragon number %d." % i, "seed": i, **base, **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 variants.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)