Borrowed from https://huggingface.co/nomic-ai/nomic-embed-text-v1.5#transformers
(input_ids: torch.LongTensor, attention_mask: torch.LongTensor, model: PreTrainedModel)
| 85 | |
| 86 | |
| 87 | def embed_prompt(input_ids: torch.LongTensor, attention_mask: torch.LongTensor, model: PreTrainedModel): |
| 88 | """ |
| 89 | Borrowed from https://huggingface.co/nomic-ai/nomic-embed-text-v1.5#transformers |
| 90 | """ |
| 91 | |
| 92 | def mean_pooling(model_output, attention_mask): |
| 93 | token_embeddings = model_output[0] |
| 94 | input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() |
| 95 | return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) |
| 96 | |
| 97 | with torch.no_grad(): |
| 98 | model_output = model(input_ids=input_ids, attention_mask=attention_mask) |
| 99 | embeddings = mean_pooling(model_output, attention_mask) |
| 100 | |
| 101 | matryoshka_dim = 512 |
| 102 | # normalize embeddings |
| 103 | embeddings = F.normalize(embeddings, p=2, dim=1) |
| 104 | embeddings = F.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],)) |
| 105 | embeddings = embeddings[:, :matryoshka_dim] |
| 106 | |
| 107 | return embeddings |
| 108 | |
| 109 | |
| 110 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected