File size: 3,201 Bytes
02bc7b8 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | from typing import Dict
import numpy as np
from attributedict.collections import AttributeDict
gen_params = AttributeDict(
[
("seed", -1), # RNG seed
("n_threads", 1), # TODO: fix this
("n_predict", 512), # new tokens to predict
("n_parts", -1), # amount of model parts (-1: determine from model dimensions)
("n_ctx", 512), # context size
("n_batch", 512), # batch size for prompt processing (must be >=32 to use BLAS)
("n_keep", 0), # number of tokens to keep from initial prompt
("n_vocab", 50272), # vocabulary size
# sampling parameters
("logit_bias", dict()), # logit bias for specific tokens: <int, float>
("top_k", 50), # <= 0 to use vocab size
("top_p", 0.95), # 1.0 = disabled
("tfs_z", 1.00), # 1.0 = disabled
("typical_p", 1.00), # 1.0 = disabled
("temp", 0.20), # 1.0 = disabled
("repeat_penalty", 1.10), # 1.0 = disabled
(
"repeat_last_n",
64,
), # last n tokens to penalize (0 = disable penalty, -1 = context size)
("frequency_penalty", 0.00), # 0.0 = disabled
("presence_penalty", 0.00), # 0.0 = disabled
("mirostat", 0), # 0 = disabled, 1 = mirostat, 2 = mirostat 2.0
("mirostat_tau", 5.00), # target entropy
("mirostat_eta", 0.10), # learning rate
]
)
class TimeStats:
def __init__(self):
self.total_tokens = 0
self.context_tokens = 0
self.context_time = 0.0
self.generation_tokens = 0
self.generation_time_list = []
def update(self, timing: Dict):
self.context_tokens = timing["context_tokens"]
self.context_time = timing["context_time"]
self.total_tokens = timing["total_tokens"]
self.generation_time_list = timing["generation_time_list"]
self.generation_tokens = len(self.generation_time_list)
self.average_speed = (self.context_time + np.sum(self.generation_time_list)) / (
self.context_tokens + self.generation_tokens
)
def show(self):
if self.total_tokens == 0:
# No stats to show.
return
print("*" * 50)
print(
f"Speed of Generation : {np.average(self.generation_time_list)*1000:.3f} ms/token"
)
print("*" * 50)
def stream_output(output_stream, time_stats: TimeStats = None):
pre = 0
for outputs in output_stream:
output_text = outputs["text"]
output_text = output_text.strip().split(" ")
now = len(output_text) - 1
if now > pre:
print(" ".join(output_text[pre:now]), end=" ", flush=True)
pre = now
print(" ".join(output_text[pre:]), flush=True)
if "timing" in outputs and outputs["timing"] is not None:
timing = outputs["timing"]
total_tokens = timing["total_tokens"]
if time_stats is not None:
time_stats.update(timing)
prompt_tokens = timing["context_tokens"]
print("-" * 50)
print("TTFT: {:.3f} s for {} tokens.".format(timing["context_time"], prompt_tokens))
return " ".join(output_text), total_tokens
|