Unverified Commit c5bb0ebd authored by Li Wang's avatar Li Wang Committed by GitHub
Browse files

[Doc] Fix prompt embedding examples (#18350)


Signed-off-by: default avatarwangli <wangli858794774@gmail.com>
parent d637b960
...@@ -36,6 +36,7 @@ llm = LLM(model=model_name, enable_prompt_embeds=True) ...@@ -36,6 +36,7 @@ llm = LLM(model=model_name, enable_prompt_embeds=True)
chat = [{"role": "user", "content": "Please tell me about the capital of France."}] chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt') token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt')
embedding_layer = transformers_model.get_input_embeddings()
prompt_embeds = embedding_layer(token_ids).squeeze(0) prompt_embeds = embedding_layer(token_ids).squeeze(0)
# Single prompt inference # Single prompt inference
...@@ -116,6 +117,7 @@ transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_nam ...@@ -116,6 +117,7 @@ transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_nam
chat = [{"role": "user", "content": "Please tell me about the capital of France."}] chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt') token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt')
embedding_layer = transformers_model.get_input_embeddings()
prompt_embeds = embedding_layer(token_ids).squeeze(0) prompt_embeds = embedding_layer(token_ids).squeeze(0)
# Prompt embeddings # Prompt embeddings
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment