Rename the class and add a parameter, 'dimension' to support both versions
Browse files- model_api.py +5 -3
model_api.py
CHANGED
@@ -6,17 +6,18 @@ from mteb import DRESModel
|
|
6 |
from tqdm import tqdm
|
7 |
|
8 |
|
9 |
-
class
|
10 |
-
def __init__(self, url: str, instruction: Optional[str] = None, batch_size: int = 128, **kwargs) -> None:
|
11 |
self.url = url
|
12 |
self.instruction = instruction
|
13 |
self.batch_size = batch_size
|
|
|
14 |
|
15 |
def get_embeddings(self, queries: List[str]) -> np.ndarray:
|
16 |
return np.asarray(
|
17 |
requests.post(self.url, json={'inputs': queries}).json()['embedding'],
|
18 |
dtype=np.float32,
|
19 |
-
).reshape(len(queries),
|
20 |
|
21 |
def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
|
22 |
return self.encode([f'{self.instruction}{query}' for query in queries])
|
@@ -38,3 +39,4 @@ class ModelV1(DRESModel):
|
|
38 |
],
|
39 |
axis=0,
|
40 |
)
|
|
|
|
6 |
from tqdm import tqdm
|
7 |
|
8 |
|
9 |
+
class SionicEmbeddingModel(DRESModel):
|
10 |
+
def __init__(self, url: str, instruction: Optional[str] = None, batch_size: int = 128, dimension: int = 2048, **kwargs) -> None:
|
11 |
self.url = url
|
12 |
self.instruction = instruction
|
13 |
self.batch_size = batch_size
|
14 |
+
self.dimension = dimension
|
15 |
|
16 |
def get_embeddings(self, queries: List[str]) -> np.ndarray:
|
17 |
return np.asarray(
|
18 |
requests.post(self.url, json={'inputs': queries}).json()['embedding'],
|
19 |
dtype=np.float32,
|
20 |
+
).reshape(len(queries), self.dimension)
|
21 |
|
22 |
def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
|
23 |
return self.encode([f'{self.instruction}{query}' for query in queries])
|
|
|
39 |
],
|
40 |
axis=0,
|
41 |
)
|
42 |
+
|