improve speed when using union of corpora

Author: Sewon Min

npm/dstore.py

@@ -547,7 +547,7 @@ def get_embs(self, indices):
         assert len(indices)==len(self.dstores)
         embs = []
         for dstore, _indices in zip(self.dstores, indices):
-            embs.append(np.array(dstore.get_embs(_indices).tolist()))
-        return np.concatenate(embs, -2)
+            embs.append(dstore.get_embs(_indices))
+        return embs
 
 

npm/npm_single.py

@@ -34,9 +34,15 @@ def decode(self, ids):
 
     def get_scores(self, queries, x):
         if type(queries)==np.ndarray:
-            all_scores = np.inner(queries, x).squeeze(1) / np.sqrt(self.dstore.dimension)
+            if type(x)==list:
+                all_scores = np.concatenate([self.get_scores(queries, xi) for xi in x], -1)
+            else:
+                all_scores = np.inner(queries, x).squeeze(1) / np.sqrt(self.dstore.dimension)
         else:
-            all_scores = torch.inner(queries, x).squeeze(1) / np.sqrt(self.dstore.dimension)
+            if type(x)==list:
+                all_scores = torch.cat([self.get_scores(queries, xi) for xi in x], -1)
+            else:
+                all_scores = torch.inner(queries, x).squeeze(1) / np.sqrt(self.dstore.dimension)
         return all_scores
 
     def get_all_scores(self, queries):