Model-Optimizer/examples/speculative_decoding/scripts/send_conversation_vllm.py

# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Send conversations from a dataset to an OpenAI-compatible endpoint."""

import argparse
import asyncio
import json
from pathlib import Path

import httpx
import openai
from openai import AsyncOpenAI
from tqdm import tqdm
from transformers import AutoTokenizer


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="""Collect hidden states from conversations
        by sending full conversations as prompts to an OpenAI-compatible endpoint."""
    )

    ## Model & Generation Parameters ##
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="Name of the served model.",
    )
    parser.add_argument(
        "--model_card",
        type=str,
        default="moonshotai/Kimi-K2-Thinking",
        help="Name of the served model card.",
    )
    parser.add_argument(
        "--meta-channel-id",
        type=str,
        default=None,
        help=(
            "Unique identifier for the meta file name used to communicate with "
            "the local serving engine. This should match the value used by the server."
        ),
    )
    ## Client Parameters ##
    parser.add_argument(
        "--base-url",
        type=str,
        default="http://localhost:8000/v1",
        help="""HTTP URL for the OpenAI-compatible endpoint.
        Defaults to `http://localhost:8000/v1`.""",
    )
    parser.add_argument(
        "--openai-api-key",
        default="EMPTY",
        help="""Access key required by the OpenAI Python client
        (not required for local serving engines like vLLM).""",
    )

    ## I/O Parameters ##
    parser.add_argument(
        "--max-seq-len",
        type=int,
        default=8192,
        help="""Maximum number of tokens in a conversation. Longer conversations will be skipped.
        Defaults to 3072 tokens.""",
    )
    parser.add_argument(
        "--input-file",
        type=Path,
        required=True,
        help="""Path to the input `jsonl` file containing conversations.
        Each entry must have a unique `conversation_id` field and a `conversations` field
        containing a list of messages.""",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        required=True,
        help="""Root directory in which to save the hidden states.
        The data will be saved as a torch (`.pt`) dump file for each conversation.""",
    )
    parser.add_argument(
        "--debug-max-num-conversations",
        type=int,
        default=None,
        help="""For debugging purposes, limit the number of conversations processed.
        Default is None, meaning no limit.""",
    )
    parser.add_argument("--num-shards", type=int, default=1, help="number of shards.")
    parser.add_argument("--shard-id-begin", type=int, default=0, help="the shard id to start.")
    parser.add_argument(
        "--shard-id-step", type=int, default=1, help="the step that the shard id progress."
    )

    return parser.parse_args()


async def main(args: argparse.Namespace) -> None:
    for shard_id in range(args.shard_id_begin, args.num_shards, args.shard_id_step):
        if args.num_shards > 1:
            input_file_path = args.input_file / "train-{:05}-{:05}.jsonl".format(
                shard_id + 1, args.num_shards
            )
        else:
            input_file_path = args.input_file

        if not input_file_path.exists():
            print(f"Input jsonl file {input_file_path} not found, skipping.")
            continue

        all_conversations = []
        with input_file_path.open("r", encoding="utf-8") as f:
            all_conversations.extend([json.loads(line) for line in f if line.strip()])

        print("Loaded", len(all_conversations), "conversations from", input_file_path)

        client: AsyncOpenAI = AsyncOpenAI(
            api_key=args.openai_api_key,
            base_url=args.base_url,
        )

        tokenizer = AutoTokenizer.from_pretrained(args.model_card, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        bos_token_id = tokenizer.bos_token_id
        if bos_token_id is None:
            raise ValueError("The tokenizer does not have a BOS token.")  # diff here

        # temp_meta_file = Path(f"/tmp/meta_{args.meta_channel_id}.json") # diff here

        temp_meta_file = Path("/dump_tmp/meta.json")  # diff here
        # INSERT_YOUR_CODE
        temp_meta_file.parent.mkdir(parents=True, exist_ok=True)
        if temp_meta_file.exists():
            print(f"Temporary meta file {temp_meta_file} found, removing it.")
            temp_meta_file.unlink()

        output_dir = args.output_dir / f"train-{(shard_id + 1):05}-{args.num_shards:05}"
        output_dir.mkdir(parents=True, exist_ok=True)
        num_invalid = 0
        num_exists = 0
        num_error = 0
        num_too_long = 0
        num_success = 0
        num_total_conversations = min(
            len(all_conversations), args.debug_max_num_conversations or len(all_conversations)
        )
        for idx, entry in enumerate(
            tqdm(
                all_conversations[: args.debug_max_num_conversations],
                desc="Processing conversations",
                total=num_total_conversations,
            )
        ):
            conversation_id = entry.get("conversation_id", None)
            if conversation_id is None:
                conversation_id = entry.get("uuid", None)
            if conversation_id is None:
                conversation_id = "{:08d}".format(idx)
            conversations = entry["conversations"]
            if not conversations or not isinstance(conversations, list):
                num_invalid += 1
                continue

            hidden_states_file = output_dir / f"{conversation_id}.pt"
            if hidden_states_file.exists():
                num_exists += 1
                continue

            # Use /tmp/meta.json to communicate with the local serving engine.
            # See usage guide for more details
            with temp_meta_file.open("w") as f:
                json.dump(
                    {
                        "conversation_id": conversation_id,
                        "output_file": str(hidden_states_file),
                    },
                    f,
                )

            input_ids = tokenizer.apply_chat_template(
                conversations, return_tensors=None, add_generation_template=False, tokenize=True
            )
            num_input_tokens = len(input_ids)
            if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len:
                num_too_long += 1
                continue
            if bos_token_id is not None and input_ids[0] == bos_token_id:
                # Remove the leading BOS token. vLLM's completion generation
                # endpoint will prepend the BOS token automatically.
                input_ids = input_ids[1:]
            input_string = tokenizer.decode(input_ids, skip_special_tokens=False)

            try:
                # Send the message to the OpenAI-compatible endpoint
                await client.completions.create(
                    model=args.model,
                    prompt=input_string,  # diff here
                    temperature=0.0,
                    max_tokens=1,
                )
            except httpx.HTTPStatusError as e:
                print(f"HTTP error for conversation {conversation_id}: {e}")
                num_error += 1
                continue
            except openai.BadRequestError:
                # Most likely the conversation is too long, ignore
                num_too_long += 1
                continue
            except Exception as e:
                num_error += 1
                print(f"Error sending conversation {conversation_id}: {e}")
                continue
            finally:
                # Ensure the meta file is cleaned up after each request
                if temp_meta_file.exists():
                    temp_meta_file.unlink()
            num_success += 1
            continue

        if num_invalid > 0:
            print(f"Skipped {num_invalid} invalid conversations without proper fields.")
        if num_too_long > 0:
            print(f"Skipped {num_too_long} conversations likely due to length constraints.")
        if num_error > 0:
            print(f"Encountered errors for {num_error} conversations.")
        if num_exists > 0:
            print(f"Skipped {num_exists} conversations that already exist.")

        if num_success == num_total_conversations:
            print(f"Successfully processed all {num_success} conversations.")
        else:
            print(
                f"Successfully processed {num_success} out of {num_total_conversations} conversations."
            )


if __name__ == "__main__":
    cli_args = parse_args()
    asyncio.run(main(cli_args))