[Model] Support Gemma2 embedding model (#9004)

2024-10-04 23:57:05 -07:00
parent 53b3a33027
commit 15986f598c
5 changed files with 99 additions and 3 deletions
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -83,6 +83,7 @@ _GENERATION_MODELS = {
 _EMBEDDING_MODELS = {
    "MistralModel": ("llama_embedding", "LlamaEmbeddingModel"),
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
+    "Gemma2Model": ("gemma2_embedding", "Gemma2EmbeddingModel"),
 }

 _MULTIMODAL_MODELS = {