Skip to content

Commit c5bb0eb

Browse files
authored
[Doc] Fix prompt embedding examples (#18350)
Signed-off-by: wangli <wangli858794774@gmail.com>
1 parent d637b96 commit c5bb0eb

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

docs/source/serving/prompt_embeds.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ llm = LLM(model=model_name, enable_prompt_embeds=True)
3636
chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
3737
token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt')
3838

39+
embedding_layer = transformers_model.get_input_embeddings()
3940
prompt_embeds = embedding_layer(token_ids).squeeze(0)
4041

4142
# Single prompt inference
@@ -116,6 +117,7 @@ transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_nam
116117
chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
117118
token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt')
118119

120+
embedding_layer = transformers_model.get_input_embeddings()
119121
prompt_embeds = embedding_layer(token_ids).squeeze(0)
120122

121123
# Prompt embeddings

0 commit comments

Comments
 (0)