HelioAI's picture
LFM3-1.7B inference modules: experimental release
67ebcef verified
Raw History Blame Contribute Delete
968 Bytes
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)