trending/Model-Optimizer
1
0
Fork 0
mirror of https://github.com/NVIDIA/Model-Optimizer.git synced 2026-10-02 03:14:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Optimize calibrate_draft_vocab to read only required lines when calib… (#618)

Browse Source 
Optimize calibrate_draft_vocab to read only required lines when
calibrate_size is set

## What does this PR do?

**Type of change:** Performance improvement

**Overview:** 
This PR optimizes the
[calibrate_draft_vocab.py](cci:7://file:///Users/obenshoham/PycharmProjects/TensorRT-Model-Optimizer/examples/speculative_decoding/scripts/calibrate_draft_vocab.py:0:0-0:0)
script to improve memory efficiency and I/O performance when using the
`--calibrate_size` parameter. Previously, the script would read all
lines from the data file into memory before slicing to the specified
`calibrate_size`, causing unnecessary resource usage for large datasets.
The optimization uses `itertools.islice` to read only the required
number of lines when `calibrate_size` is specified.

## Usage

The script usage remains unchanged. When using `--calibrate_size`, the
script now only reads the specified number of lines instead of loading
the entire dataset:

```bash
# Only reads first 1000 lines from the dataset (optimized)
python scripts/calibrate_draft_vocab.py \
    --model meta-llama/Llama-3.2-1B-Instruct \
    --data input_conversations/daring-anteater.jsonl \
    --draft_vocab_size 32000 \
    --calibrate_size 1000 \
    --save_dir draft_vocab_cache

Signed-off-by: Ofir Ben Shoham <ofir_benshoham@intuit.com>
This commit is contained in:
Ben Shoham Ofir
2025-12-19 17:47:08 +00:00
committed by GitHub
parent bdd10c2dbe
commit 2a51bbd7cb
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/speculative_decoding/scripts/calibrate_draft_vocab.py
+3 -3
View File
@@ -16,6 +16,7 @@
import argparse
import json
import os
from itertools import islice
import torch
from transformers import AutoTokenizer
@@ -47,9 +48,8 @@ def main():
print("Calibrating vocab...")
tokenizer = AutoTokenizer.from_pretrained(args.model)
with open(args.data) as f:
conversations = [json.loads(line)["conversations"] for line in f]
if args.calibrate_size:
conversations = conversations[: args.calibrate_size]
lines = islice(f, args.calibrate_size) if args.calibrate_size else f
conversations = [json.loads(line)["conversations"] for line in lines]
conversations = [item for sublist in conversations for item in sublist]
d2t = calibrate_frequent_vocab(tokenizer, conversations, args.draft_vocab_size)