| 104 | |
| 105 | |
| 106 | def send_prompt(data: dict) -> tuple[float, list[float]]: |
| 107 | session = data["session"] |
| 108 | server_address: str = data["server_address"] |
| 109 | |
| 110 | t_submit = time() |
| 111 | if data["external_server"]: |
| 112 | json_data: dict = { |
| 113 | "prompt": data["prompt"], "ignore_eos": True, |
| 114 | "seed": data["seed"], "max_tokens": data["n_predict"], "stream": True} |
| 115 | response = session.post(f"{server_address}/v1/completions", json=json_data, stream=True) |
| 116 | elif data["synthetic_prompt"]: |
| 117 | json_data: dict = { |
| 118 | "prompt": data["prompt"], "ignore_eos": True, "cache_prompt": False, |
| 119 | "seed": data["seed"], "n_predict": data["n_predict"], "stream": True} |
| 120 | response = session.post(f"{server_address}/completion", json=json_data, stream=True) |
| 121 | else: |
| 122 | response = session.post( |
| 123 | f"{server_address}/apply-template", |
| 124 | json={"messages": [{"role": "user", "content": data["prompt"], "stream": True}]} |
| 125 | ) |
| 126 | response.raise_for_status() |
| 127 | prompt: str = json.loads(response.text)["prompt"] |
| 128 | |
| 129 | json_data: dict = {"prompt": prompt, "seed": data["seed"], "n_predict": data["n_predict"], "stream": True} |
| 130 | response = session.post(f"{server_address}/completion", json=json_data, stream=True) |
| 131 | response.raise_for_status() |
| 132 | |
| 133 | lines = [] |
| 134 | token_arrival_times: list[float] = [] |
| 135 | for line in response.iter_lines(decode_unicode=False): |
| 136 | if not line.startswith(b"data: "): |
| 137 | continue |
| 138 | lines.append(line) |
| 139 | token_arrival_times.append(time()) |
| 140 | token_arrival_times = token_arrival_times[:-1] |
| 141 | if len(lines) > 1 and "timings" in json.loads(lines[-2][6:]): |
| 142 | token_arrival_times = token_arrival_times[:-1] |
| 143 | |
| 144 | return (t_submit, token_arrival_times) |
| 145 | |
| 146 | |
| 147 | def benchmark( |