universal-dependencies pipeline

README.md

@@ -60,4 +60,12 @@ nlp=UDgoeswith("KoichiYasuoka/roberta-classical-chinese-base-ud-goeswith")
 print(nlp("孟子見梁惠王"))
 ```
 
+or
+
+```
+from transformers import pipeline
+nlp=pipeline("universal-dependencies","KoichiYasuoka/roberta-classical-chinese-base-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": 0,
   "classifier_dropout": null,
+  "custom_pipelines": {
+    "universal-dependencies": {
+      "impl": "ud.UniversalDependenciesPipeline"
+    }
+  },
   "eos_token_id": 2,
   "gradient_checkpointing": false,
   "hidden_act": "gelu",

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)