Spaces:
Running on Zero
Running on Zero
File size: 1,096 Bytes
8ef9ed7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | """Sentence Transformers adapter used for both documents and queries."""
from __future__ import annotations
class SentenceTransformerEncoder:
def __init__(self, model_id: str = "magibu/embeddingmagibu-200m", device: str | None = None):
from sentence_transformers import SentenceTransformer
self.model_id = model_id
self.model = SentenceTransformer(model_id, device=device, tokenizer_kwargs={"extra_special_tokens": {}})
def token_count(self, text: str) -> int:
return len(self.model.tokenizer.encode(text, add_special_tokens=True))
def encode_documents(self, texts: list[str]) -> list[list[float]]:
encoder = getattr(self.model, "encode_document", self.model.encode)
return encoder(texts, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False).tolist()
def encode_queries(self, texts: list[str]) -> list[list[float]]:
encoder = getattr(self.model, "encode_query", self.model.encode)
return encoder(texts, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False).tolist()
|