improve goeswith

README.md

@@ -39,6 +39,12 @@ class UDgoeswith(object):
       e=self.model(input_ids=torch.tensor(x)).logits.numpy()[:,1:-2,:]
     r=[1 if i==0 else -1 if j.endswith("|root") else 0 for i,j in sorted(self.model.config.id2label.items())]
     e+=numpy.where(numpy.add.outer(numpy.identity(e.shape[0]),r)==0,0,numpy.nan)
+    g=self.model.config.label2id["X|_|goeswith"]
+    r=numpy.tri(e.shape[0])
+    for i in range(e.shape[0]):
+      for j in range(i+2,e.shape[1]):
+        r[i,j]=r[i,j-1] if numpy.nanargmax(e[i,j-1])==g else 1
+    e[:,:,g]+=numpy.where(r==0,0,numpy.nan)
     m=numpy.full((e.shape[0]+1,e.shape[1]+1),numpy.nan)
     m[1:,1:]=numpy.nanmax(e,axis=2).transpose()
     p=numpy.zeros(m.shape)

ud.py

@@ -13,6 +13,12 @@ class UniversalDependenciesPipeline(TokenClassificationPipeline):
     e=model_output["logits"].numpy()
     r=[1 if i==0 else -1 if j.endswith("|root") else 0 for i,j in sorted(self.model.config.id2label.items())]
     e+=numpy.where(numpy.add.outer(numpy.identity(e.shape[0]),r)==0,0,numpy.nan)
+    g=self.model.config.label2id["X|_|goeswith"]
+    r=numpy.tri(e.shape[0])
+    for i in range(e.shape[0]):
+      for j in range(i+2,e.shape[1]):
+        r[i,j]=r[i,j-1] if numpy.nanargmax(e[i,j-1])==g else 1
+    e[:,:,g]+=numpy.where(r==0,0,numpy.nan)
     m=numpy.full((e.shape[0]+1,e.shape[1]+1),numpy.nan)
     m[1:,1:]=numpy.nanmax(e,axis=2).transpose()
     p=numpy.zeros(m.shape)