Put all the VOCAB_PREFIX_SIZE fiddling in the same place

Author: AngledLuffa

stanza/models/depparse/model.py

@@ -283,6 +283,7 @@ def pack(x):
         else:
             loss = 0
             preds.append(F.log_softmax(unlabeled_scores, 2).detach().cpu().numpy())
-            preds.append(deprel_scores.max(3)[1].detach().cpu().numpy())
+            deprels = deprel_scores.max(3)[1].detach().cpu().numpy() + VOCAB_PREFIX_SIZE
+            preds.append(deprels)
 
         return loss, preds

stanza/models/depparse/trainer.py

@@ -18,7 +18,6 @@
 from stanza.models.common.foundation_cache import load_bert, load_bert_with_peft, NoTransformerFoundationCache
 from stanza.models.common.chuliu_edmonds import chuliu_edmonds_one_root
 from stanza.models.common.peft_config import build_peft_wrapper, load_peft_wrapper
-from stanza.models.common.vocab import VOCAB_PREFIX_SIZE
 from stanza.models.depparse.model import Parser
 from stanza.models.pos.vocab import MultiVocab
 
@@ -150,7 +149,6 @@ def predict(self, batch, unsort=True):
         self.model.eval()
         batch_size = word.size(0)
         _, preds = self.model(word, word_mask, wordchars, wordchars_mask, upos, xpos, ufeats, pretrained, lemma, head, deprel, word_orig_idx, sentlens, wordlens, text)
-        preds[1] += VOCAB_PREFIX_SIZE
         head_seqs = [chuliu_edmonds_one_root(adj[:l, :l])[1:] for adj, l in zip(preds[0], sentlens)] # remove attachment for the root
         deprel_seqs = [self.vocab['deprel'].unmap([preds[1][i][j+1][h] for j, h in enumerate(hs)]) for i, hs in enumerate(head_seqs)]