algorithm improved

ud.py

@@ -7,6 +7,8 @@ class UniversalDependenciesPipeline(TokenClassificationPipeline):
     with torch.no_grad():
       e=self.model(input_ids=torch.tensor([v[0:i]+[self.tokenizer.mask_token_id]+v[i+1:]+[j] for i,j in enumerate(v[1:-1],1)],device=self.device))
     return {"logits":e.logits[:,1:-2,:],**model_inputs}
+  def check_model_type(self,supported_models):
+    pass
   def postprocess(self,model_outputs,**kwargs):
     import numpy
     if "logits" not in model_outputs:
@@ -29,17 +31,34 @@ class UniversalDependenciesPipeline(TokenClassificationPipeline):
       h=self.chu_liu_edmonds(m)
     v=[(s,e) for s,e in model_outputs["offset_mapping"][0].tolist() if s<e]
     q=[self.model.config.id2label[p[j,i]].split("|") for i,j in enumerate(h)]
-    g="aggregation_strategy" in kwargs and kwargs["aggregation_strategy"]!="none"
-    if g:
+    if "aggregation_strategy" in kwargs and kwargs["aggregation_strategy"]!="none":
       for i,j in reversed(list(enumerate(q[1:],1))):
         if j[-1]=="goeswith" and set([t[-1] for t in q[h[i]+1:i+1]])=={"goeswith"}:
           h=[b if i>b else b-1 for a,b in enumerate(h) if i!=a]
           v[i-1]=(v[i-1][0],v.pop(i)[1])
           q.pop(i)
+        elif v[i-1][1]>v[i][0]:
+          h=[b if i>b else b-1 for a,b in enumerate(h) if i!=a]
+          v[i-1]=(v[i-1][0],v.pop(i)[1])
+          q.pop(i)
     t=model_outputs["sentence"].replace("\n"," ")
+    for i,(s,e) in reversed(list(enumerate(v))):
+      w=t[s:e]
+      if w.startswith(" "):
+        j=len(w)-len(w.lstrip())
+        w=w.lstrip()
+        v[i]=(v[i][0]+j,v[i][1])
+      if w.endswith(" "):
+        j=len(w)-len(w.rstrip())
+        w=w.rstrip()
+        v[i]=(v[i][0],v[i][1]-j)
+      if w.strip()=="":
+        h=[b if i>b else b-1 for a,b in enumerate(h) if i!=a]
+        v.pop(i)
+        q.pop(i)
     u="# text = "+t+"\n"
     for i,(s,e) in enumerate(v):
-      u+="\t".join([str(i+1),t[s:e],t[s:e] if g else "_",q[i][0],"_","|".join(q[i][1:-1]),str(0 if h[i]==i else h[i]+1),q[i][-1],"_","_" if i+1<len(v) and e<v[i+1][0] else "SpaceAfter=No"])+"\n"
+      u+="\t".join([str(i+1),t[s:e],"_",q[i][0],"_","|".join(q[i][1:-1]),str(0 if h[i]==i else h[i]+1),q[i][-1],"_","_" if i+1<len(v) and e<v[i+1][0] else "SpaceAfter=No"])+"\n"
     return u+"\n"
   def chu_liu_edmonds(self,matrix):
     import numpy