File tree Expand file tree Collapse file tree 1 file changed +2
-0
lines changed Expand file tree Collapse file tree 1 file changed +2
-0
lines changed Original file line number Diff line number Diff line change @@ -36,6 +36,7 @@ llm = LLM(model=model_name, enable_prompt_embeds=True)
3636chat = [{" role" : " user" , " content" : " Please tell me about the capital of France." }]
3737token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt = True , return_tensors = ' pt' )
3838
39+ embedding_layer = transformers_model.get_input_embeddings()
3940prompt_embeds = embedding_layer(token_ids).squeeze(0 )
4041
4142# Single prompt inference
@@ -116,6 +117,7 @@ transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_nam
116117chat = [{" role" : " user" , " content" : " Please tell me about the capital of France." }]
117118token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt = True , return_tensors = ' pt' )
118119
120+ embedding_layer = transformers_model.get_input_embeddings()
119121prompt_embeds = embedding_layer(token_ids).squeeze(0 )
120122
121123# Prompt embeddings
You can’t perform that action at this time.
0 commit comments