@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 158 lines 5.08 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """vLLM OpenAI-Compatible Client with Prompt Embeddings.
4
5 This script demonstrates how to:
6 1. Generate prompt embeddings using Hugging Face Transformers.
7 2. Encode them in base64 format.
8 3. Send them to a vLLM server for inference via both:
9 - OpenAI-compatible Chat Completions API
10 - OpenAI-compatible Completions API
11
12 Important distinction between the two APIs:
13
14 - Chat Completions API: `prompt_embeds` content parts should encode ONLY
15 the user-provided content, not a templated conversation. The server
16 renders the surrounding chat template around the embedded content at
17 request time, the same way it would for a plain text `content` string.
18 Embedding a full templated conversation here would double-apply the
19 template and likely produce undesirable results.
20
21 - Completions API: the server does NOT apply a chat template to
22 `prompt_embeds`. The caller is responsible for producing embeddings for
23 the full, already-templated prompt (i.e. apply the chat template first,
24 then embed the resulting token IDs). Anything the model would normally
25 need (system prompt, role markers, generation prompt, etc.) must already
26 be baked into the embedded tokens.
27
28 Run the vLLM server first:
29 vllm serve meta-llama/Llama-3.2-1B-Instruct \
30 --runner generate \
31 --max-model-len 4096 \
32 --enable-prompt-embeds
33
34 Run the client:
35 python examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py
36
37 Model: meta-llama/Llama-3.2-1B-Instruct
38 Note: This model is gated on Hugging Face Hub.
39 You must request access to use it:
40 https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct
41
42 Dependencies:
43 - transformers
44 - torch
45 - openai
46 """
47
48 import transformers
49 from openai import OpenAI
50
51 from vllm.utils.serial_utils import tensor2base64
52
53
54 def run_chat_completion_prompt_embeds(
55 client: OpenAI,
56 model_name: str,
57 tokenizer: transformers.PreTrainedTokenizerBase,
58 embedding_layer,
59 messages: list[dict],
60 ) -> None:
61 """Run a Chat Completions API request using prompt_embeds content parts.
62
63 This example embeds ONLY the user-provided content of the final user turn, the
64 vLLM server applies the chat template around it at request time.
65 """
66 user_content = messages[-1]["content"]
67 content_token_ids = tokenizer(
68 user_content, return_tensors="pt", add_special_tokens=False
69 ).input_ids
70 content_prompt_embeds = embedding_layer(content_token_ids).squeeze(0)
71 encoded_embeds = tensor2base64(content_prompt_embeds)
72
73 api_messages = [
74 *messages[:-1],
75 {
76 "role": messages[-1]["role"],
77 "content": [{"type": "prompt_embeds", "data": encoded_embeds}],
78 },
79 ]
80
81 chat_completion = client.chat.completions.create(
82 model=model_name,
83 max_tokens=6,
84 temperature=0.0,
85 messages=api_messages,
86 )
87
88 print("-" * 30)
89 print("Chat Completions API")
90 print(chat_completion.choices[0].message.content)
91 print("-" * 30)
92
93
94 def run_completion_prompt_embeds(
95 client: OpenAI,
96 model_name: str,
97 tokenizer: transformers.PreTrainedTokenizerBase,
98 embedding_layer,
99 messages: list[dict],
100 ) -> None:
101 """Run a Completions API request using prompt embeddings.
102
103 The Completions endpoint does not apply a chat template,
104 so the caller must apply it and embed the full templated prompt.
105 """
106 templated_token_ids = tokenizer.apply_chat_template(
107 messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
108 ).input_ids
109 templated_prompt_embeds = embedding_layer(templated_token_ids).squeeze(0)
110 encoded_embeds = tensor2base64(templated_prompt_embeds)
111
112 completion = client.completions.create(
113 model=model_name,
114 prompt=None,
115 max_tokens=6,
116 temperature=0.0,
117 # NOTE: The OpenAI client allows passing in extra JSON body via the
118 # `extra_body` argument.
119 extra_body={"prompt_embeds": encoded_embeds},
120 )
121
122 print("-" * 30)
123 print("Completions API")
124 print(completion.choices[0].text)
125 print("-" * 30)
126
127
128 def main() -> None:
129 client = OpenAI(
130 api_key="EMPTY",
131 base_url="http://localhost:8000/v1",
132 )
133
134 model_name = "meta-llama/Llama-3.2-1B-Instruct"
135
136 tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
137 transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
138 embedding_layer = transformers_model.get_input_embeddings()
139
140 messages = [
141 {"role": "user", "content": "Please tell me about the capital of France."}
142 ]
143
144 # Chat Completions API: embed ONLY the user content. The server wraps
145 # the embedding in the chat template when it renders the messages.
146 run_chat_completion_prompt_embeds(
147 client, model_name, tokenizer, embedding_layer, messages
148 )
149
150 # Completions API: embed the FULL templated prompt. The caller must
151 # apply the chat template up-front.
152 run_completion_prompt_embeds(
153 client, model_name, tokenizer, embedding_layer, messages
154 )
155
156
157 if __name__ == "__main__":
158 main()