algorithm improved

ud.py

@@ -7,10 +7,13 @@ class BellmanFordTokenClassificationPipeline(TokenClassificationPipeline):
     x=self.model.config.label2id
     y=[k for k in x if k.find("|")<0 and not k.startswith("I-")]
     self.transition=numpy.full((len(x),len(x)),-numpy.inf)
+    self.ilabel=numpy.full(len(x),-numpy.inf)
     for k,v in x.items():
       if k.find("|")<0:
         for j in ["I-"+k[2:]] if k.startswith("B-") else [k]+y if k.startswith("I-") else y:
           self.transition[v,x[j]]=0
+        if k.startswith("I-"):
+          self.ilabel[v]=0
   def check_model_type(self,supported_models):
     pass
   def postprocess(self,model_outputs,**kwargs):
@@ -19,6 +22,10 @@ class BellmanFordTokenClassificationPipeline(TokenClassificationPipeline):
     return self.bellman_ford_token_classification(model_outputs,**kwargs)
   def bellman_ford_token_classification(self,model_outputs,**kwargs):
     m=model_outputs["logits"][0].numpy()
+    x=model_outputs["offset_mapping"][0].tolist()
+    for i,(s,e) in enumerate(x):
+      if i>0 and s<e and x[i-1][1]>s:
+        m[i]+=self.ilabel
     e=numpy.exp(m-numpy.max(m,axis=-1,keepdims=True))
     z=e/e.sum(axis=-1,keepdims=True)
     for i in range(m.shape[0]-1,0,-1):
@@ -26,13 +33,16 @@ class BellmanFordTokenClassificationPipeline(TokenClassificationPipeline):
     k=[numpy.argmax(m[0]+self.transition[0])]
     for i in range(1,m.shape[0]):
       k.append(numpy.argmax(m[i]+self.transition[k[-1]]))
-    w=[{"entity":self.model.config.id2label[j],"start":s,"end":e,"score":z[i,j]} for i,((s,e),j) in enumerate(zip(model_outputs["offset_mapping"][0].tolist(),k)) if s<e]
+    w=[{"entity":self.model.config.id2label[j],"start":s,"end":e,"score":z[i,j]} for i,((s,e),j) in enumerate(zip(x,k)) if s<e]
     if "aggregation_strategy" in kwargs and kwargs["aggregation_strategy"]!="none":
       for i,t in reversed(list(enumerate(w))):
         p=t.pop("entity")
         if p.startswith("I-"):
           w[i-1]["score"]=min(w[i-1]["score"],t["score"])
           w[i-1]["end"]=w.pop(i)["end"]
+        elif i>0 and w[i-1]["end"]>t["start"]:
+          w[i-1]["score"]=min(w[i-1]["score"],t["score"])
+          w[i-1]["end"]=w.pop(i)["end"]
         elif p.startswith("B-"):
           t["entity_group"]=p[2:]
         else: