cached_file support

setup.py

@@ -5,7 +5,7 @@ URL="https://github.com/KoichiYasuoka/SuPar-Kanbun"
 
 setuptools.setup(
   name="suparkanbun",
-  version="1.4.2",
+  version="1.4.6",
   description="Tokenizer POS-tagger and Dependency-parser for Classical Chinese",
   long_description=long_description,
   long_description_content_type="text/markdown",
@@ -16,10 +16,10 @@ setuptools.setup(
   keywords="NLP Chinese",
   packages=setuptools.find_packages(),
   install_requires=[
-    "supar>=1.1.3",
+    "supar>=1.1.4",
     "transformers>=4.0.1",
     "spacy>=2.2.2",
-    "deplacy>=2.0.2"
+    "deplacy>=2.0.3"
   ],
   python_requires=">=3.7",
   package_data={"suparkanbun":["models/*.txt","models/*/*.txt","models/*/*.json"]},

suparkanbun/download.py

@@ -3,24 +3,30 @@
 
 import os
 
-def download(url,file,dir="."):
+try:
+  from transformers.utils import cached_file
+except:
+  from transformers.file_utils import cached_path,hf_bucket_url
+  cached_file=lambda x,y:cached_path(hf_bucket_url(x,y))
+
+def download(rootdir,file,dir="."):
   import shutil
-  from transformers.file_utils import cached_path
   t=os.path.join(dir,"filesize.txt")
-  shutil.copy(cached_path(url+"filesize.txt"),t)
+  shutil.copy(cached_file(rootdir,os.path.dirname(file)+"/filesize.txt"),t)
   with open(t,"r") as f:
     r=f.read()
+  f=os.path.basename(file)
   ft=0
   for t in r.split("\n"):
     s=t.split()
     if len(s)==2:
-      if s[0]==file:
+      if s[0]==f:
         ft=int(s[1])
   if ft==0:
     return
-  shutil.copy(cached_path(url+file),os.path.join(dir,file))
+  shutil.copy(cached_file(rootdir,file),os.path.join(dir,f))
 
-def checkdownload(url,dir="."):
+def checkdownload(rootdir,model,dir="."):
   while True:
     t=os.path.join(dir,"filesize.txt")
     with open(t,"r") as f:
@@ -35,7 +41,7 @@ def checkdownload(url,dir="."):
         except:
           j=-1
         if i!=j:
-          download(url,s[0],dir)
+          download(rootdir,model+s[0],dir)
           break
     else:
       return

suparkanbun/suparkanbun.py

@@ -5,9 +5,6 @@ import os
 PACKAGE_DIR=os.path.abspath(os.path.dirname(__file__))
 DOWNLOAD_DIR=os.path.join(PACKAGE_DIR,"models")
 
-from transformers.file_utils import hf_bucket_url
-MODEL_URL=hf_bucket_url("KoichiYasuoka/SuPar-Kanbun","suparkanbun/models/")
-
 import numpy
 from spacy.language import Language
 from spacy.symbols import LANG,NORM,LEMMA,POS,TAG,DEP,HEAD
@@ -36,7 +33,7 @@ class SuParKanbunLanguage(Language):
       "name":"SuParKanbun_lzh",
       "parent_package":"suparkanbun",
       "pipeline":"Tokenizer, POS-Tagger, Parser",
-      "spacy_version":">=2.1.0"
+      "spacy_version":">=2.2.2"
     }
     self._path=None
 
@@ -59,7 +56,7 @@ class SuParKanbunTokenizer(object):
     self.supar=Parser.load(f)
     if danku:
       d=os.path.join(DOWNLOAD_DIR,bert+".danku")
-      self.danku=AutoModelTagger(d,["B","E","E2","E3","M","S"])
+      self.danku=AutoModelTagger(d,["B","E","E2","E3","M","S"],[("B","E"),("B","E2"),("B","E3"),("B","M"),("E","B"),("E","S"),("E2","E"),("E3","E2"),("M","E3"),("M","M"),("S","B"),("S","S")])
     else:
       self.danku=None
     self.gloss=MakeGloss()
@@ -167,19 +164,32 @@ class SuParKanbunTokenizer(object):
     return doc
 
 class AutoModelTagger(object):
-  def __init__(self,dir,label=None):
+  def __init__(self,dir,label=None,links=None):
     from suparkanbun.download import checkdownload
     from transformers import AutoModelForTokenClassification,AutoTokenizer
-    checkdownload(MODEL_URL+os.path.basename(dir)+"/",dir)
+    import numpy
+    checkdownload("KoichiYasuoka/SuPar-Kanbun","suparkanbun/models/"+os.path.basename(dir)+"/",dir)
     self.model=AutoModelForTokenClassification.from_pretrained(dir)
     self.tokenizer=AutoTokenizer.from_pretrained(dir)
     self.label=label if label else self.model.config.id2label
+    if links:
+      self.transition=numpy.full((len(self.label),len(self.label)),numpy.nan)
+      x=self.model.config.label2id
+      for f,t in links:
+        self.transition[x[f],x[t]]=0
+    else:
+      self.transition=numpy.zeros((len(self.label),len(self.label)))
   def __call__(self,text):
-    import torch
-    input=self.tokenizer.encode(text,return_tensors="pt")
-    output=self.model(input)
-    predict=torch.argmax(output[0],dim=2)
-    return [(t,self.label[p]) for t,p in zip(text,predict[0].tolist()[1:])]
+    import torch,numpy
+    v=self.tokenizer(text,return_offsets_mapping=True)
+    with torch.no_grad():
+      m=self.model(torch.tensor([v["input_ids"]])).logits[0].numpy()
+    for i in range(m.shape[0]-1,0,-1):
+      m[i-1]+=numpy.nanmax(m[i]+self.transition,axis=1)
+    p=[numpy.nanargmax(m[0])]
+    for i in range(1,m.shape[0]):
+      p.append(numpy.nanargmax(m[i]+self.transition[p[-1]]))
+    return [(text[t[0]:t[1]],self.label[q]) for t,q in zip(v["offset_mapping"],p) if t[0]<t[1]]
 
 class MakeGloss(object):
   def __init__(self,file=None):