algorithm improved

ud.py

@@ -19,9 +19,10 @@ class UniversalDependenciesPipeline(TokenClassificationPipeline):
       f=True
       k=r["input_ids"]
     except:
-      r=list(r)[0]
       f=False
-    w=self.tokenizer.convert_ids_to_tokens(r["input_ids"][0])
+      r=list(r)[0]
+      k=r["input_ids"]
+    w=self.tokenizer.convert_ids_to_tokens(k[0])
     if len(m)!=len(w):
       for i,j in enumerate(w):
         if j.endswith("@@"):
@@ -30,9 +31,7 @@ class UniversalDependenciesPipeline(TokenClassificationPipeline):
           m[i]=(s,s+len(j)-2)
     r["offset_mapping"]=torch.tensor([m]).to(self.device)
     r["sentence"]=sentence
-    if f:
-      return r
-    return iter([r])
+    return r if f else iter([r])
   def _forward(self,model_inputs):
     import torch
     v=model_inputs["input_ids"][0].tolist()