Support Alibaba-NLP/gte-Qwen2-7B-instruct embedding Model (#1186)

Co-authored-by: Ying Sheng <sqy1415@gmail.com>
This commit is contained in:
Chayenne
2024-08-25 10:29:12 -07:00
committed by GitHub
co-authored by Ying Sheng
parent 66e7dcaf70
commit 30b4f771b0
15 changed files with 167 additions and 55 deletions
@@ -29,7 +29,11 @@ class LlamaEmbeddingModel(nn.Module):
positions: torch.Tensor,
input_metadata: InputMetadata,
input_embeds: torch.Tensor = None,
get_embedding: bool = True,
) -> EmbeddingPoolerOutput:
assert (
get_embedding
), "LlamaEmbeddingModel / MistralModel is only used for embedding"
hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
return self.pooler(hidden_states, input_metadata)