File size: 1,497 Bytes
f54cf3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
from typing import Dict, List, Optional, Union
import numpy as np
import requests
from mteb import DRESModel
from tqdm import tqdm
class SionicEmbeddingModel(DRESModel):
def __init__(self, url: str, instruction: Optional[str] = None, batch_size: int = 128, dimension: int = 2048, **kwargs) -> None:
self.url = url
self.instruction = instruction
self.batch_size = batch_size
self.dimension = dimension
def get_embeddings(self, queries: List[str]) -> np.ndarray:
return np.asarray(
requests.post(self.url, json={'inputs': queries}).json()['embedding'],
dtype=np.float32,
).reshape(len(queries), self.dimension)
def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
return self.encode([f'{self.instruction}{query}' for query in queries])
def encode_corpus(self, corpus: List[Union[Dict[str, str], str]], **kwargs) -> np.ndarray:
sentences: List[str] = (
[f"{doc.get('title', '')} {doc['text']}".strip() for doc in corpus]
if isinstance(corpus[0], dict)
else corpus
)
return self.encode(sentences)
def encode(self, sentences: List[str], **kwargs) -> np.ndarray:
return np.concatenate(
[
self.get_embeddings(sentences[idx:idx + self.batch_size])
for idx in tqdm(range(0, len(sentences), self.batch_size), desc='encode')
],
axis=0,
)
|