sionic commited on
Commit
18a7736
1 Parent(s): f36e968

Rename the class and add a parameter, 'dimension' to support both versions

Browse files
Files changed (1) hide show
  1. 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 ModelV1(DRESModel):
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), 2048)
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
+