| | from pylate import models |
| | from transformers import AutoTokenizer |
| | import torch |
| | import numpy as np |
| |
|
| | class EndpointHandler: |
| | def __init__(self, path=""): |
| | self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True) |
| | self.model = models.ColBERT(model_name_or_path=path) |
| | self.model.eval() |
| |
|
| | def _to_list(self, emb): |
| | """ |
| | Make the output JSON-serialisable: |
| | – torch.Tensor ➜ emb.cpu().tolist() |
| | – np.ndarray ➜ emb.tolist() |
| | – list[...] ➜ recurse |
| | """ |
| | if isinstance(emb, torch.Tensor): |
| | return emb.cpu().tolist() |
| | if isinstance(emb, np.ndarray): |
| | return emb.tolist() |
| | if isinstance(emb, list): |
| | return [self._to_list(e) for e in emb] |
| | return emb |
| |
|
| | def __call__(self, data): |
| | texts = data.get("inputs") or data.get("text") or data |
| | if isinstance(texts, str): |
| | texts = [texts] |
| |
|
| | with torch.no_grad(): |
| | emb = self.model.encode( |
| | texts, |
| | is_query=True, |
| | batch_size=32, |
| | ) |
| |
|
| | return self._to_list(emb) |
| |
|
