(texts: list[str], tokenizer, embed_model)
| 696 | |
| 697 | |
| 698 | async def hf_embedding(texts: list[str], tokenizer, embed_model) -> np.ndarray: |
| 699 | device = next(embed_model.parameters()).device |
| 700 | input_ids = tokenizer( |
| 701 | texts, return_tensors="pt", padding=True, truncation=True |
| 702 | ).input_ids.to(device) |
| 703 | with torch.no_grad(): |
| 704 | outputs = embed_model(input_ids) |
| 705 | embeddings = outputs.last_hidden_state.mean(dim=1) |
| 706 | if embeddings.dtype == torch.bfloat16: |
| 707 | return embeddings.detach().to(torch.float32).cpu().numpy() |
| 708 | else: |
| 709 | return embeddings.detach().cpu().numpy() |
| 710 | |
| 711 | |
| 712 | async def ollama_embedding(texts: list[str], embed_model, **kwargs) -> np.ndarray: |
no outgoing calls
no test coverage detected