use numpy.nan

ud.py

@@ -22,8 +22,8 @@ class UniversalDependenciesPipeline(TokenClassificationPipeline):
     h=self.chu_liu_edmonds(m)
     z=[i for i,j in enumerate(h) if i==j]
     if len(z)>1:
-      k,h=z[numpy.nanargmax(m[z,z])],numpy.nanmin(m)-numpy.nanmax(m)
-      m[:,z]+=[[0 if j in z and (i!=j or i==k) else h for i in z] for j in range(m.shape[0])]
+      k=z[numpy.nanargmax(m[z,z])]
+      m[:,z]+=[[0 if j in z and (i!=j or i==k) else numpy.nan for i in z] for j in range(m.shape[0])]
       h=self.chu_liu_edmonds(m)
     v=[(s,e) for s,e in model_output["offset_mapping"][0].tolist() if s<e]
     q=[self.model.config.id2label[p[j,i]].split("|") for i,j in enumerate(h)]
@@ -53,9 +53,9 @@ class UniversalDependenciesPipeline(TokenClassificationPipeline):
         return h
     y,x=[i for i,j in enumerate(x) if j==max(x)],[i for i,j in enumerate(x) if j<max(x)]
     z=matrix-numpy.nanmax(matrix,axis=0)
-    m=numpy.block([[z[x,:][:,x],numpy.max(z[x,:][:,y],axis=1).reshape(len(x),1)],[numpy.max(z[y,:][:,x],axis=0),numpy.max(z[y,y])]])
-    k=[j if i==len(x) else x[j] if j<len(x) else y[numpy.argmax(z[y,x[i]])] for i,j in enumerate(self.chu_liu_edmonds(m))]
+    m=numpy.block([[z[x,:][:,x],numpy.nanmax(z[x,:][:,y],axis=1).reshape(len(x),1)],[numpy.nanmax(z[y,:][:,x],axis=0),numpy.nanmax(z[y,y])]])
+    k=[j if i==len(x) else x[j] if j<len(x) else y[numpy.nanargmax(z[y,x[i]])] for i,j in enumerate(self.chu_liu_edmonds(m))]
     h=[j if i in y else k[x.index(i)] for i,j in enumerate(h)]
-    i=y[numpy.argmax(z[k[-1],y] if k[-1]<len(x) else z[y,y])]
+    i=y[numpy.nanargmax(z[k[-1],y] if k[-1]<len(x) else z[y,y])]
     h[i]=x[k[-1]] if k[-1]<len(x) else i
     return h