universal-dependencies pipeline

README.md

@@ -57,4 +57,12 @@ nlp=UDgoeswith("KoichiYasuoka/roberta-large-japanese-aozora-ud-goeswith")
 print(nlp("全学年にわたって小学校の国語の教科書に挿し絵が用いられている"))
 ```
 
+or
+
+```
+from transformers import pipeline
+nlp=pipeline("universal-dependencies","KoichiYasuoka/roberta-large-japanese-aozora-ud-goeswith",trust_remote_code=True)
+print(nlp("全学年にわたって小学校の国語の教科書に挿し絵が用いられている"))
+```
+
 [ufal.chu-liu-edmonds](https://pypi.org/project/ufal.chu-liu-edmonds/) is required.

config.json

@@ -5,6 +5,11 @@
   "attention_probs_dropout_prob": 0.1,
   "bos_token_id": 2,
   "classifier_dropout": null,
+  "custom_pipelines": {
+    "universal-dependencies": {
+      "impl": "ud.UniversalDependenciesPipeline"
+    }
+  },
   "eos_token_id": 3,
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,

ud.py

@@ -0,0 +1,35 @@
+from transformers import TokenClassificationPipeline
+from transformers.pipelines import PIPELINE_REGISTRY
+
+class UniversalDependenciesPipeline(TokenClassificationPipeline):
+  def _forward(self,model_input):
+    import torch
+    v=model_input["input_ids"][0].tolist()
+    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)]))
+    return {"logits":e.logits[:,1:-2,:],**model_input}
+  def postprocess(self,model_output,**kwargs):
+    import numpy
+    import ufal.chu_liu_edmonds
+    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)
+    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)
+    p[1:,1:]=numpy.nanargmax(e,axis=2).transpose()
+    for i in range(1,m.shape[0]):
+      m[i,0],m[i,i],p[i,0]=m[i,i],numpy.nan,p[i,i]
+    h=ufal.chu_liu_edmonds.chu_liu_edmonds(m)[0]
+    if [0 for i in h if i==0]!=[0]:
+      m[:,0]+=numpy.where(m[:,0]<numpy.nanmax(m[:,0]),numpy.nan,0)
+      h=ufal.chu_liu_edmonds.chu_liu_edmonds(m)[0]
+    t=model_output["sentence"]
+    u="# text = "+t+"\n"
+    v=[(s,e) for s,e in model_output["offset_mapping"][0].tolist() if s<e]
+    for i,(s,e) in enumerate(v,1):
+      q=self.model.config.id2label[p[i,h[i]]].split("|")
+      u+="\t".join([str(i),t[s:e],"_",q[0],"_","|".join(q[1:-1]),str(h[i]),q[-1],"_","_" if i<len(v) and e<v[i][0] else "SpaceAfter=No"])+"\n"
+    return u+"\n"
+
+PIPELINE_REGISTRY.register_pipeline("universal-dependencies",pipeline_class=UniversalDependenciesPipeline)