TransformersUD improved

README.md

@@ -65,16 +65,17 @@ class TransformersUD(object):
     w=[(t["start"],t["end"],t["entity_group"]) for t in self.deprel(text)]
     z,n={t["start"]:t["entity"].split("|") for t in self.tagger(text)},len(w)
     r,m=[text[s:e] for s,e,p in w],numpy.full((n+1,n+1),numpy.nan)
-    v=self.tokenizer(r,add_special_tokens=False)["input_ids"]
+    v,c=self.tokenizer(r,add_special_tokens=False)["input_ids"],[]
     for i,t in enumerate(v):
       q=[self.tokenizer.cls_token_id]+t+[self.tokenizer.sep_token_id]
-      c=[q]+v[0:i]+[[self.tokenizer.mask_token_id]]+v[i+1:]+[[q[-1]]]
-      b=[len(sum(c[0:j+1],[])) for j in range(len(c))]
-      d=self.model(input_ids=torch.tensor([sum(c,[])]),
-        token_type_ids=torch.tensor([[0]*b[0]+[1]*(b[-1]-b[0])]))
-      s,e=d.start_logits.tolist()[0],d.end_logits.tolist()[0]
+      c.append([q]+v[0:i]+[[self.tokenizer.mask_token_id]]+v[i+1:]+[[q[-1]]])
+    b=[[len(sum(x[0:j+1],[])) for j in range(len(x))] for x in c]
+    d=self.model(input_ids=torch.tensor([sum(x,[]) for x in c]),
+      token_type_ids=torch.tensor([[0]*x[0]+[1]*(x[-1]-x[0]) for x in b]))
+    s,e=d.start_logits.tolist(),d.end_logits.tolist()
+    for i in range(n):
       for j in range(n):
-        m[i+1,0 if i==j else j+1]=s[b[j]]+e[b[j+1]-1]
+        m[i+1,0 if i==j else j+1]=s[i][b[i][j]]+e[i][b[i][j+1]-1]
     i=numpy.nanargmax(m[:,0])
     m[0:i,0]=m[i+1:,0]=numpy.nan
     h=ufal.chu_liu_edmonds.chu_liu_edmonds(m)[0]