Upload model

modeling_gzipembed.py

@@ -24,7 +24,7 @@ class GZIPEmbeddingModel(PreTrainedModel):
             ncd = [0] * len(self.config.corpus)
             with multiprocessing.Pool(num_procs) as pool:
                 data = enumerate(self.config.corpus)
-                results = pool.map(self.ncd_r,data)
+                results = pool.map(self.ncd_r,(data,p))
             for i,row in results:
                 ncd[i]=row
             x.append(ncd)
@@ -33,10 +33,10 @@ class GZIPEmbeddingModel(PreTrainedModel):
             x = x.to(self.reduction_head.dtype).to(self.reduction_head.device)
             return self.reduction_head(x)
         return x if not return_tensor else torch.tensor(x)
-        
-    def ncd_r(self,r):
+
+    def ncd_r(self,rp):
         i=r[0]
-        return i,self.ncd(r[1],p)
+        return i,self.ncd(r[1],r[2])
 
     def normalize(self, x):
         x = ''.join([char for char in x.lower() if char in "abcdefghijklmnopqrstuvwxyz "])