| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from statistics import mean |
| 5 | from typing import Any, NamedTuple |
| 6 | |
| 7 | import numpy as np # type: ignore |
| 8 | import pandas as pd # type: ignore |
| 9 | from bench_utils import ( |
| 10 | TEXT_SEPARATOR, |
| 11 | Color, |
| 12 | logger, |
| 13 | ) |
| 14 | from tqdm import tqdm |
| 15 | from transformers import AutoTokenizer # type: ignore |
| 16 | |
| 17 | # Conversation ID is a string (e.g: "UzTK34D") |
| 18 | ConvId = str |
| 19 | |
| 20 | # A list of dicts (dicts with keys "id" and "messages") |
| 21 | ShareGptConversations = list[dict[str, Any]] |
| 22 | |
| 23 | # A list of dicts (dicts with keys "role" and "content") |
| 24 | MessagesList = list[dict[str, str]] |
| 25 | |
| 26 | # Map conversation ID to conversation messages |
| 27 | ConversationsMap = list[ConvId, MessagesList] |
| 28 | |
| 29 | |
| 30 | class Distribution(ABC): |
| 31 | @abstractmethod |
| 32 | def sample(self, size: int = 1) -> np.ndarray: |
| 33 | pass |
| 34 | |
| 35 | |
| 36 | class UniformDistribution(Distribution): |
| 37 | def __init__( |
| 38 | self, |
| 39 | min_val: int | float, |
| 40 | max_val: int | float, |
| 41 | is_integer: bool = True, |
| 42 | ) -> None: |
| 43 | self.min_val = min_val |
| 44 | self.max_val = max_val |
| 45 | self.is_integer = is_integer |
| 46 | |
| 47 | def sample(self, size: int = 1) -> np.ndarray: |
| 48 | if self.is_integer: |
| 49 | return np.random.randint( |
| 50 | int(self.min_val), int(self.max_val + 1), size=size |
| 51 | ) |
| 52 | else: |
| 53 | return np.random.uniform(self.min_val, self.max_val, size=size) |
| 54 | |
| 55 | def __repr__(self) -> str: |
| 56 | return f"UniformDistribution[{self.min_val}, {self.max_val}]" |
| 57 | |
| 58 | |
| 59 | class ConstantDistribution(Distribution): |
| 60 | def __init__(self, value: int | float) -> None: |
| 61 | self.value = value |
| 62 | self.max_val = value |
| 63 | |
| 64 | def sample(self, size: int = 1) -> np.ndarray: |
| 65 | return np.full(shape=size, fill_value=self.value) |
| 66 | |
| 67 | def __repr__(self) -> str: |
| 68 | return f"Constant[{self.value}]" |
| 69 | |
| 70 | |
| 71 | class ZipfDistribution(Distribution): |
| 72 | def __init__(self, alpha: float, max_val: int | None = None) -> None: |
| 73 | self.alpha = alpha |
| 74 | self.max_val = max_val |
| 75 | |
| 76 | def sample(self, size: int = 1) -> np.ndarray: |
| 77 | samples = np.random.zipf(self.alpha, size=size) |
| 78 | if self.max_val: |
| 79 | samples = np.minimum(samples, self.max_val) |
| 80 | return samples |
| 81 | |
| 82 | def __repr__(self) -> str: |
| 83 | return f"ZipfDistribution[{self.alpha}]" |
| 84 | |
| 85 | |
| 86 | class PoissonDistribution(Distribution): |
| 87 | def __init__(self, alpha: float, max_val: int | None = None) -> None: |
| 88 | self.alpha = alpha |
| 89 | self.max_val = max_val |
| 90 | |
| 91 | def sample(self, size: int = 1) -> np.ndarray: |
| 92 | samples = np.random.poisson(self.alpha, size=size) |
| 93 | if self.max_val: |
| 94 | samples = np.minimum(samples, self.max_val) |
| 95 | return samples |
| 96 | |
| 97 | def __repr__(self) -> str: |
| 98 | return f"PoissonDistribution[{self.alpha}]" |
| 99 | |
| 100 | |
| 101 | class LognormalDistribution(Distribution): |
| 102 | def __init__( |
| 103 | self, |
| 104 | mean: float | None = None, |
| 105 | sigma: float | None = None, |
| 106 | average: int | None = None, |
| 107 | median_ratio: float | None = None, |
| 108 | max_val: int | None = None, |
| 109 | ) -> None: |
| 110 | self.average = average |
| 111 | self.median_ratio = median_ratio |
| 112 | self.max_val = max_val |
| 113 | |
| 114 | if average is not None: |
| 115 | if average < 1: |
| 116 | raise ValueError("Lognormal average must be positive") |
| 117 | |
| 118 | if mean or sigma: |
| 119 | raise ValueError( |
| 120 | "When using lognormal average, you can't provide mean/sigma" |
| 121 | ) |
| 122 | |
| 123 | if self.median_ratio is None: |
| 124 | # Default value that provides relatively wide range of values |
| 125 | self.median_ratio = 0.85 |
| 126 | |
| 127 | # Calculate mean/sigma of np.random.lognormal based on the average |
| 128 | mean, sigma = self._generate_lognormal_by_median( |
| 129 | target_average=self.average, median_ratio=self.median_ratio |
| 130 | ) |
| 131 | else: |
| 132 | if mean is None or sigma is None: |
| 133 | raise ValueError( |
| 134 | "Must provide both mean and sigma if average is not used" |
| 135 | ) |
| 136 | |
| 137 | if mean <= 0 or sigma < 0: |
| 138 | raise ValueError( |
| 139 | "Lognormal mean must be positive and sigma must be non-negative" |
| 140 | ) |
| 141 | |
| 142 | # Mean and standard deviation of the underlying normal distribution |
| 143 | # Based on numpy.random.lognormal |
| 144 | self.mean = mean |
| 145 | self.sigma = sigma |
| 146 | |
| 147 | @staticmethod |
| 148 | def _generate_lognormal_by_median( |
| 149 | target_average: int, median_ratio: float |
| 150 | ) -> tuple[float, float]: |
| 151 | """Compute (mu, sigma) for a lognormal distribution given: |
| 152 | - a target average (mean of the distribution) |
| 153 | - a ratio of median / mean (controls skewness), assume mean > median |
| 154 | |
| 155 | Background: |
| 156 | If Z ~ Normal(mu, sigma^2), then X = exp(Z) ~ LogNormal(mu, sigma). |
| 157 | * mean(X) = exp(mu + sigma^2 / 2) |
| 158 | * median(X) = exp(mu) |
| 159 | |
| 160 | So: |
| 161 | median / mean = exp(mu) / exp(mu + sigma^2 / 2) |
| 162 | = exp(-sigma^2 / 2) |
| 163 | |
| 164 | Rearranging: |
| 165 | sigma^2 = 2 * ln(mean / median) |
| 166 | mu = ln(median) |
| 167 | |
| 168 | This gives a unique (mu, sigma) for any valid mean and median. |
| 169 | """ |
| 170 | # Check input validity: median must be smaller than mean |
| 171 | if median_ratio <= 0 or median_ratio >= 1: |
| 172 | raise ValueError("median_ratio must be in range (0, 1)") |
| 173 | |
| 174 | target_median = target_average * median_ratio |
| 175 | |
| 176 | # Solve sigma^2 = 2 * ln(mean / median) |
| 177 | sigma = np.sqrt(2 * np.log(target_average / target_median)) |
| 178 | mu = np.log(target_median) |
| 179 | |
| 180 | return mu, sigma |
| 181 | |
| 182 | def sample(self, size: int = 1) -> np.ndarray: |
| 183 | samples = np.random.lognormal(mean=self.mean, sigma=self.sigma, size=size) |
| 184 | |
| 185 | if self.average is not None: |
| 186 | # Scale to average |
| 187 | samples *= self.average / samples.mean() |
| 188 | |
| 189 | if self.max_val: |
| 190 | samples = np.minimum(samples, self.max_val) |
| 191 | |
| 192 | return np.round(samples).astype(int) |
| 193 | |
| 194 | def __repr__(self) -> str: |
| 195 | if self.average: |
| 196 | return ( |
| 197 | f"LognormalDistribution[{self.average}, " |
| 198 | f"{self.median_ratio}, {self.max_val}]" |
| 199 | ) |
| 200 | return f"LognormalDistribution[{self.mean}, {self.sigma}, {self.max_val}]" |
| 201 | |
| 202 | |
| 203 | class GenConvArgs(NamedTuple): |
| 204 | num_conversations: int |
| 205 | text_files: list[str] |
| 206 | input_num_turns: Distribution |
| 207 | input_common_prefix_num_tokens: Distribution |
| 208 | input_prefix_num_tokens: Distribution |
| 209 | input_num_tokens: Distribution |
| 210 | output_num_tokens: Distribution |
| 211 | print_stats: bool |
| 212 | |
| 213 | |
| 214 | def verify_field_exists( |
| 215 | conf: dict, field_name: str, section: str, subsection: str |
| 216 | ) -> None: |
| 217 | if field_name not in conf: |
| 218 | raise ValueError( |
| 219 | f"Missing field '{field_name}' in {section=} and {subsection=}" |
| 220 | ) |
| 221 | |
| 222 | |
| 223 | def get_random_distribution( |
| 224 | conf: dict, section: str, subsection: str, optional: bool = False |
| 225 | ) -> Distribution: |
| 226 | # section can be "prompt_input" or "prompt_output" (both required) |
| 227 | conf = conf[section] |
| 228 | |
| 229 | if optional and subsection not in conf: |
| 230 | # Optional subsection, if not found assume the value is always 0 |
| 231 | return ConstantDistribution(0) |
| 232 | |
| 233 | # subsection can be "num_turns", "num_tokens" or "prefix_num_tokens" |
| 234 | if subsection not in conf: |
| 235 | raise ValueError(f"Missing subsection {subsection} in section {section}") |
| 236 | |
| 237 | conf = conf[subsection] |
| 238 | |
| 239 | distribution = conf.get("distribution") |
| 240 | if distribution is None: |
| 241 | raise ValueError( |
| 242 | f"Missing field 'distribution' in {section=} and {subsection=}" |
| 243 | ) |
| 244 | |
| 245 | if distribution == "constant": |
| 246 | verify_field_exists(conf, "value", section, subsection) |
| 247 | return ConstantDistribution(conf["value"]) |
| 248 | |
| 249 | elif distribution == "zipf": |
| 250 | verify_field_exists(conf, "alpha", section, subsection) |
| 251 | max_val = conf.get("max", None) |
| 252 | return ZipfDistribution(conf["alpha"], max_val=max_val) |
| 253 | |
| 254 | elif distribution == "poisson": |
| 255 | verify_field_exists(conf, "alpha", section, subsection) |
| 256 | max_val = conf.get("max", None) |
| 257 | return PoissonDistribution(conf["alpha"], max_val=max_val) |
| 258 | |
| 259 | elif distribution == "lognormal": |
| 260 | max_val = conf.get("max", None) |
| 261 | |
| 262 | if "average" in conf: |
| 263 | # Infer lognormal mean/sigma (numpy) from input average |
| 264 | median_ratio = conf.get("median_ratio", None) |
| 265 | return LognormalDistribution( |
| 266 | average=conf["average"], median_ratio=median_ratio, max_val=max_val |
| 267 | ) |
| 268 | |
| 269 | # Use mean/sigma directly (for full control over the distribution) |
| 270 | verify_field_exists(conf, "mean", section, subsection) |
| 271 | verify_field_exists(conf, "sigma", section, subsection) |
| 272 | return LognormalDistribution( |
| 273 | mean=conf["mean"], sigma=conf["sigma"], max_val=max_val |
| 274 | ) |
| 275 | |
| 276 | elif distribution == "uniform": |
| 277 | verify_field_exists(conf, "min", section, subsection) |
| 278 | verify_field_exists(conf, "max", section, subsection) |
| 279 | |
| 280 | min_value = conf["min"] |
| 281 | max_value = conf["max"] |
| 282 | |
| 283 | assert min_value > 0 |
| 284 | assert min_value <= max_value |
| 285 | |
| 286 | is_integer = isinstance(min_value, int) and isinstance(max_value, int) |
| 287 | return UniformDistribution(min_value, max_value, is_integer) |
| 288 | else: |
| 289 | raise ValueError(f"Unknown distribution: {distribution}") |
| 290 | |
| 291 | |
| 292 | def parse_input_json_file(conf: dict) -> GenConvArgs: |
| 293 | # Validate the input file |
| 294 | assert isinstance(conf, dict) |
| 295 | required_fields = [ |
| 296 | "filetype", |
| 297 | "num_conversations", |
| 298 | "text_files", |
| 299 | "prompt_input", |
| 300 | "prompt_output", |
| 301 | ] |
| 302 | for field in required_fields: |
| 303 | assert field in conf, f"Missing field {field} in input {conf}" |
| 304 | |
| 305 | assert conf["filetype"] == "generate_conversations" |
| 306 | |
| 307 | assert conf["num_conversations"] > 0, "num_conversations should be larger than zero" |
| 308 | |
| 309 | text_files = conf["text_files"] |
| 310 | |
| 311 | assert isinstance(text_files, list), "Field 'text_files' should be a list" |
| 312 | assert len(text_files) > 0, ( |
| 313 | "Field 'text_files' should be a list with at least one file" |
| 314 | ) |
| 315 | |
| 316 | # Parse the parameters for the prompt input/output workload |
| 317 | input_num_turns = get_random_distribution(conf, "prompt_input", "num_turns") |
| 318 | input_num_tokens = get_random_distribution(conf, "prompt_input", "num_tokens") |
| 319 | input_common_prefix_num_tokens = get_random_distribution( |
| 320 | conf, "prompt_input", "common_prefix_num_tokens", optional=True |
| 321 | ) |
| 322 | input_prefix_num_tokens = get_random_distribution( |
| 323 | conf, "prompt_input", "prefix_num_tokens" |
| 324 | ) |
| 325 | output_num_tokens = get_random_distribution(conf, "prompt_output", "num_tokens") |
| 326 | |
| 327 | print_stats: bool = conf.get("print_stats", False) |
| 328 | assert isinstance(print_stats, bool), ( |
| 329 | "Field 'print_stats' should be either 'true' or 'false'" |
| 330 | ) |
| 331 | |
| 332 | args = GenConvArgs( |
| 333 | num_conversations=conf["num_conversations"], |
| 334 | text_files=text_files, |
| 335 | input_num_turns=input_num_turns, |
| 336 | input_common_prefix_num_tokens=input_common_prefix_num_tokens, |
| 337 | input_prefix_num_tokens=input_prefix_num_tokens, |
| 338 | input_num_tokens=input_num_tokens, |
| 339 | output_num_tokens=output_num_tokens, |
| 340 | print_stats=print_stats, |
| 341 | ) |
| 342 | return args |
| 343 | |
| 344 | |
| 345 | def print_conv_stats(conversations: ConversationsMap, tokenizer: AutoTokenizer) -> None: |
| 346 | # Collect statistics |
| 347 | conv_stats: list[dict[Any, Any]] = [] |
| 348 | req_stats: list[int] = [] |
| 349 | |
| 350 | print("\nCollecting statistics...") |
| 351 | for messages in conversations.values(): |
| 352 | # messages is a list of dicts |
| 353 | user_tokens: list[int] = [] |
| 354 | assistant_tokens: list[int] = [] |
| 355 | request_tokens: list[int] = [] |
| 356 | |
| 357 | req_tokens = 0 |
| 358 | for m in messages: |
| 359 | content = m["content"] |
| 360 | num_tokens = len(tokenizer(content).input_ids) |
| 361 | |
| 362 | if m["role"] == "user": |
| 363 | user_tokens.append(num_tokens) |
| 364 | # New user prompt including all chat history |
| 365 | req_tokens += num_tokens |
| 366 | request_tokens.append(req_tokens) |
| 367 | |
| 368 | elif m["role"] == "assistant": |
| 369 | assistant_tokens.append(num_tokens) |
| 370 | # Update assistant answer |
| 371 | # (will be part of chat history for the next user prompt) |
| 372 | req_tokens += num_tokens |
| 373 | |
| 374 | item_stats = { |
| 375 | "conversation_turns": len(messages), |
| 376 | "user_tokens": mean(user_tokens), |
| 377 | "assistant_tokens": mean(assistant_tokens), |
| 378 | } |
| 379 | |
| 380 | conv_stats.append(item_stats) |
| 381 | req_stats.extend(request_tokens) |
| 382 | |
| 383 | # Print statistics |
| 384 | percentiles = [0.25, 0.5, 0.75, 0.9, 0.99] |
| 385 | |
| 386 | print(TEXT_SEPARATOR) |
| 387 | print(f"{Color.YELLOW}Conversations statistics:{Color.RESET}") |
| 388 | print(TEXT_SEPARATOR) |
| 389 | df = pd.DataFrame(conv_stats) |
| 390 | print(df.describe(percentiles=percentiles).transpose()) |
| 391 | print(TEXT_SEPARATOR) |
| 392 | print(f"{Color.YELLOW}Request statistics:{Color.RESET}") |
| 393 | print(TEXT_SEPARATOR) |
| 394 | df = pd.DataFrame(req_stats, columns=["request_tokens"]) |
| 395 | print(df.describe(percentiles=percentiles).transpose()) |
| 396 | print(TEXT_SEPARATOR) |
| 397 | |
| 398 | |
| 399 | def generate_conversations( |
| 400 | args: GenConvArgs, tokenizer: AutoTokenizer |
| 401 | ) -> ConversationsMap: |
| 402 | # Text for all user prompts |
| 403 | # (text from the input text files will be appended to this line) |
| 404 | base_prompt_text = "Please rewrite the following text and add more content: " |
| 405 | base_prompt_token_count = len( |
| 406 | tokenizer.encode(base_prompt_text, add_special_tokens=False) |
| 407 | ) |
| 408 | |
| 409 | logger.info(f"{Color.PURPLE}Generating conversations...{Color.RESET}") |
| 410 | logger.info(args) |
| 411 | |
| 412 | list_of_tokens = [] |
| 413 | |
| 414 | for filename in args.text_files: |
| 415 | # Load text file that will be used to generate prompts |
| 416 | with open(filename) as file: |
| 417 | data = file.read() |
| 418 | tokens_in_file = tokenizer.encode(data, add_special_tokens=False) |
| 419 | list_of_tokens.extend(tokens_in_file) |
| 420 | logger.info( |
| 421 | f"Loaded {len(tokens_in_file)} tokens from file {filename}, " |
| 422 | f"total tokens so far: {len(list_of_tokens)}" |
| 423 | ) |
| 424 | |
| 425 | conversations: ConversationsMap = {} |
| 426 | conv_id = 0 |
| 427 | |
| 428 | # Generate number of turns for every conversation |
| 429 | turn_count: np.ndarray = args.input_num_turns.sample(args.num_conversations) |
| 430 | |
| 431 | # Turn count should be at least 2 (one user prompt and one assistant answer) |
| 432 | turn_count = np.maximum(turn_count, 2) |
| 433 | |
| 434 | # Round up to an even number (every user prompt should have an answer) |
| 435 | turn_count = turn_count + (turn_count % 2) |
| 436 | |
| 437 | # Generate number of prefix tokens for every conversation |
| 438 | conv_prefix_tokens: np.ndarray = args.input_prefix_num_tokens.sample( |
| 439 | args.num_conversations |
| 440 | ) |
| 441 | |
| 442 | # Used to reduce shared text between conversations |
| 443 | # (jump/skip over text sections between conversations) |
| 444 | base_offset = 0 |
| 445 | |
| 446 | # Common prefix size for all conversations (only 1 sample required) |
| 447 | common_prefix_text = "" |
| 448 | common_prefix_tokens: int = args.input_common_prefix_num_tokens.sample(1)[0] |
| 449 | if common_prefix_tokens > 0: |
| 450 | # Using "." at the end to separate sentences |
| 451 | common_prefix_text = ( |
| 452 | tokenizer.decode(list_of_tokens[: common_prefix_tokens - 2]) + "." |
| 453 | ) |
| 454 | base_offset += common_prefix_tokens |
| 455 | |
| 456 | for conv_id in tqdm( |
| 457 | range(args.num_conversations), |
| 458 | total=args.num_conversations, |
| 459 | desc="Generating conversations", |
| 460 | unit="conv", |
| 461 | ): |
| 462 | # Generate a single conversation |
| 463 | messages: MessagesList = [] |
| 464 | |
| 465 | nturns = turn_count[conv_id] |
| 466 | |
| 467 | # User prompt token count per turn (with lower limit) |
| 468 | input_token_count: np.ndarray = args.input_num_tokens.sample(nturns).astype(int) |
| 469 | input_token_count = np.maximum(input_token_count, base_prompt_token_count) |
| 470 | |
| 471 | # Assistant answer token count per turn (with lower limit) |
| 472 | output_token_count: np.ndarray = args.output_num_tokens.sample(nturns).astype( |
| 473 | int |
| 474 | ) |
| 475 | output_token_count = np.maximum(output_token_count, 1) |
| 476 | |
| 477 | user_turn = True |
| 478 | for turn_id in range(nturns): |
| 479 | if user_turn: |
| 480 | role = "user" |
| 481 | num_tokens = input_token_count[turn_id] |
| 482 | |
| 483 | # Generate the user prompt, |
| 484 | # use a unique prefix (the conv_id) for each conversation |
| 485 | # (to avoid shared prefix between conversations) |
| 486 | content = f"{conv_id} is a nice number... " |
| 487 | |
| 488 | if len(common_prefix_text) > 0 and turn_id == 0: |
| 489 | content = common_prefix_text + content |
| 490 | |
| 491 | # Update the number of tokens left for the content |
| 492 | num_tokens -= len(tokenizer.encode(content, add_special_tokens=False)) |
| 493 | |
| 494 | if turn_id == 0: |
| 495 | prefix_num_tokens = conv_prefix_tokens[conv_id] |
| 496 | if prefix_num_tokens > 0: |
| 497 | # Add prefix text (context) to the first turn |
| 498 | start_offset = base_offset |
| 499 | end_offset = start_offset + prefix_num_tokens |
| 500 | assert len(list_of_tokens) > end_offset, ( |
| 501 | "Not enough input text to generate " |
| 502 | f"{prefix_num_tokens} tokens for the " |
| 503 | f"prefix text ({start_offset=}, {end_offset=})" |
| 504 | ) |
| 505 | |
| 506 | content += f"{conv_id}, " + tokenizer.decode( |
| 507 | list_of_tokens[start_offset:end_offset] |
| 508 | ) |
| 509 | base_offset += prefix_num_tokens |
| 510 | |
| 511 | # Add the actual user prompt/question after the prefix text |
| 512 | content += base_prompt_text |
| 513 | num_tokens -= base_prompt_token_count |
| 514 | |
| 515 | if num_tokens > 0: |
| 516 | # Add text from the input file (to reach the desired token count) |
| 517 | start_offset = base_offset + turn_id * input_token_count.max() |
| 518 | end_offset = start_offset + num_tokens |
| 519 | assert len(list_of_tokens) > end_offset, ( |
| 520 | f"Not enough input text to generate {num_tokens} tokens " |
| 521 | f"for the prompt ({start_offset=}, {end_offset=})" |
| 522 | ) |
| 523 | |
| 524 | # Convert tokens back to text |
| 525 | content += tokenizer.decode(list_of_tokens[start_offset:end_offset]) |
| 526 | else: |
| 527 | role = "assistant" |
| 528 | # This content will not be used as input to the LLM server |
| 529 | # (actual answers will be used instead). |
| 530 | # Content is only required to determine the min_tokens/max_tokens |
| 531 | # (inputs to the LLM server). |
| 532 | num_tokens = output_token_count[turn_id] |
| 533 | assert len(list_of_tokens) > num_tokens, ( |
| 534 | f"Not enough input text to generate {num_tokens} " |
| 535 | "tokens for assistant content" |
| 536 | ) |
| 537 | content = tokenizer.decode(list_of_tokens[:num_tokens]) |
| 538 | |
| 539 | # Append the user/assistant message to the list of messages |
| 540 | messages.append({"role": role, "content": content}) |
| 541 | user_turn = not user_turn |
| 542 | |
| 543 | # Add the new conversation |
| 544 | conversations[f"CONV_ID_{conv_id}"] = messages |
| 545 | |
| 546 | # Increase base offset for the next conversation |
| 547 | base_offset += nturns |
| 548 | |
| 549 | if args.print_stats: |
| 550 | print_conv_stats(conversations, tokenizer) |
| 551 | |
| 552 | return conversations |
| 553 | |
| 554 | |
| 555 | def conversations_list_to_dict(input_list: ShareGptConversations) -> ConversationsMap: |
| 556 | conversations: ConversationsMap = {} |
| 557 | |
| 558 | for item in input_list: |
| 559 | conv_id: str = item["id"] |
| 560 | assert isinstance(conv_id, str) |
| 561 | |
| 562 | assert conv_id not in conversations, ( |
| 563 | f"Conversation ID {conv_id} found more than once in the input" |
| 564 | ) |
| 565 | |
| 566 | messages: MessagesList = item["messages"] |
| 567 | assert isinstance(messages, list), ( |
| 568 | f"Conversation messages should be a list (ID: {conv_id})" |
| 569 | ) |
| 570 | assert len(messages) > 0, f"Conversation with no messages (ID: {conv_id})" |
| 571 | |
| 572 | conversations[conv_id] = messages |
| 573 | |
| 574 | logger.info(f"Using {len(conversations)} unique conversations (IDs)") |
| 575 | assert len(conversations) == len(input_list) |
| 576 | |
| 577 | # Print statistics about the selected conversations |
| 578 | stats: list[dict[str, Any]] = [] |
| 579 | for conv_data in conversations.values(): |
| 580 | stats.append({"num_turns": len(conv_data)}) |
| 581 | |
| 582 | print(TEXT_SEPARATOR) |
| 583 | print(f"{Color.YELLOW}Conversations statistics:{Color.RESET}") |
| 584 | print(TEXT_SEPARATOR) |
| 585 | percentiles = [0.25, 0.5, 0.75, 0.9, 0.99, 0.999, 0.9999] |
| 586 | conv_stats = pd.DataFrame(stats).describe(percentiles=percentiles) |
| 587 | print(conv_stats.transpose()) |
| 588 | print(TEXT_SEPARATOR) |
| 589 | |
| 590 | return conversations |
| 591 | |
| 592 | |
| 593 | def conversations_dict_to_list(input_dict: ConversationsMap) -> ShareGptConversations: |
| 594 | output: ShareGptConversations = [] |
| 595 | for conv_id, conv_data in input_dict.items(): |
| 596 | new_item = {"id": conv_id, "messages": conv_data} |
| 597 | output.append(new_item) |
| 598 | |
| 599 | return output |