Deal with use_vocab option

yzhangcs · yzhangcs · commit d93ca2b4599d · 2020-11-23T22:10:05.000+08:00
diff --git a/supar/utils/field.py b/supar/utils/field.py
@@ -362,5 +362,7 @@ def build(self, dataset, min_freq=1):
 
     def transform(self, charts):
         charts = [self.preprocess(chart) for chart in charts]
-        charts = [torch.tensor([[self.vocab[i] if i is not None else -1 for i in row] for row in chart]) for chart in charts]
+        if self.use_vocab:
+            charts = [[[self.vocab[i] if i is not None else -1 for i in row] for row in chart] for chart in charts]
+        charts = [torch.tensor(chart) for chart in charts]
         return charts
