TensorRT LLM customized dataset implementation. It assumes the dataset to be consist of several lines of json, each line is a minimal OpenAI API format request. Example format of each sample on each line: { "input": { "messages": [ {
| 685 | |
| 686 | |
| 687 | class CustomDataset(BenchmarkDataset): |
| 688 | """ |
| 689 | TensorRT LLM customized dataset implementation. |
| 690 | It assumes the dataset to be consist of several lines of json, each line is a minimal OpenAI API format request. |
| 691 | Example format of each sample on each line: |
| 692 | { |
| 693 | "input": { |
| 694 | "messages": [ |
| 695 | { |
| 696 | "role": "system", |
| 697 | "content": "" |
| 698 | }, |
| 699 | { |
| 700 | "role": "user", |
| 701 | "content": "" |
| 702 | } |
| 703 | ], |
| 704 | "max_tokens": 2048, |
| 705 | } |
| 706 | } |
| 707 | """ |
| 708 | |
| 709 | def __init__(self, dataset_path: str, **kwargs) -> None: |
| 710 | super().__init__(**kwargs) |
| 711 | self.dataset_path = dataset_path |
| 712 | self.data = [] |
| 713 | self.load_data() |
| 714 | |
| 715 | def load_data(self) -> None: |
| 716 | if self.dataset_path is None: |
| 717 | raise ValueError("--dataset-path is not provided") |
| 718 | with open(self.dataset_path, encoding="utf-8") as f: |
| 719 | for line in f: |
| 720 | self.data.append(json.loads(line)) |
| 721 | random.shuffle(self.data) |
| 722 | |
| 723 | def sample(self, tokenizer: PreTrainedTokenizerBase, |
| 724 | num_requests: int) -> list[SampleRequest]: |
| 725 | """ |
| 726 | Optimized version using batch tokenization for better performance. |
| 727 | """ |
| 728 | # Collect all prompts and metadata |
| 729 | prompts = [] |
| 730 | max_tokens_list = [] |
| 731 | prompt_lengths = [] |
| 732 | |
| 733 | for i, entry in enumerate(self.data): |
| 734 | if len(prompts) >= num_requests: |
| 735 | break |
| 736 | prompt = entry["input"]["messages"][1]["content"] |
| 737 | max_tokens = entry["input"]["max_tokens"] |
| 738 | prompts.append(prompt) |
| 739 | max_tokens_list.append(max_tokens) |
| 740 | if "num_tokens" in entry["input"] and isinstance( |
| 741 | entry["input"]["num_tokens"], |
| 742 | int) and entry["input"]["num_tokens"] > 0: |
| 743 | prompt_lengths.append(entry["input"]["num_tokens"]) |
| 744 |