diff --git a/inference/worker/__main__.py b/inference/worker/__main__.py index e5c15fb4..190ed788 100644 --- a/inference/worker/__main__.py +++ b/inference/worker/__main__.py @@ -41,25 +41,6 @@ def main( prompt = prefix + "\n".join(messages) + "\nAssistant:" - # TODO: use the seed - # torch.manual_seed(work_request.seed) - # model_output = pipe(prompt, max_new_tokens=work_request.max_new_tokens, do_sample=True, return_full_text=False)[ - # 0 - # ]["generated_text"] - # model_output = model_output.strip() - - # # fake streaming - # split_idcs = [m.start() for m in re.finditer(r"([\w:]+)", model_output)] - # pieces = [model_output[a:b] for a, b in zip([0] + split_idcs, split_idcs + [None])] - # for piece in pieces: - # if not piece: - # continue - # if piece.strip() in ("User:", "Assistant:"): - # break - # ws.send(inference.WorkResponsePacket(token=piece).json()) - # time.sleep(0.1) - # ws.send(inference.WorkResponsePacket(is_end=True).json()) - response = requests.post( f"{inference_server_url}/generate_stream", json={ @@ -70,6 +51,7 @@ def main( "top_k": work_request.top_k, "top_p": work_request.top_p, "temperature": work_request.temperature, + "seed": work_request.seed, }, }, stream=True, diff --git a/oasst-shared/oasst_shared/schemas/inference.py b/oasst-shared/oasst_shared/schemas/inference.py index b50cef9c..91a16b61 100644 --- a/oasst-shared/oasst_shared/schemas/inference.py +++ b/oasst-shared/oasst_shared/schemas/inference.py @@ -13,7 +13,7 @@ class WorkRequest(pydantic.BaseModel): conversation: protocol.Conversation = pydantic.Field(..., repr=False) model_name: str = "distilgpt2" max_new_tokens: int = 100 - seed: int = pydantic.Field(default_factory=lambda: random.randint(0, 2**32 - 1)) + seed: int = pydantic.Field(default_factory=lambda: random.randint(-(2**31), 2**31 - 1)) do_sample: bool = True top_k: int = 50 top_p: float = 0.9