Authenticate, embed newline-separated ``text``, and return normalised vectors.
(tokenizer, model, device, request: Request, text: str)
| 86 | |
| 87 | |
| 88 | def run_embed(tokenizer, model, device, request: Request, text: str): |
| 89 | """Authenticate, embed newline-separated ``text``, and return normalised vectors.""" |
| 90 | api_key = request.headers.get("x-api-key") |
| 91 | if api_key != os.environ["API_KEY"]: |
| 92 | raise HTTPException(status_code=401, detail="Unauthorized") |
| 93 | |
| 94 | texts = [t for t in text.split("\n") if t.strip()] |
| 95 | if not texts: |
| 96 | return [] |
| 97 | |
| 98 | print(f"Start embedding {len(texts)} texts") |
| 99 | try: |
| 100 | with torch.no_grad(): |
| 101 | batch_dict = tokenizer(texts, padding=True, truncation=True, return_tensors="pt") |
| 102 | batch_dict = {k: v.to(device) for k, v in batch_dict.items()} |
| 103 | |
| 104 | outputs = model(**batch_dict) |
| 105 | embeddings = average_pool(outputs.last_hidden_state, batch_dict["attention_mask"]) |
| 106 | embeddings = F.normalize(embeddings, p=2, dim=1) |
| 107 | embeddings = embeddings.cpu().numpy().tolist() |
| 108 | |
| 109 | print("Finished embedding texts.") |
| 110 | return embeddings |
| 111 | |
| 112 | except RuntimeError as e: |
| 113 | print(f"Error during embedding: {str(e)}") |
| 114 | if "CUDA out of memory" in str(e): |
| 115 | print("CUDA OOM. Try reducing batch size or using a smaller model.") |
| 116 | raise |
no test coverage detected