[fix] eos --> <eos> (#2224)

wenet-e2e · Dec 12, 2023 · 7facff5 · 7facff5
1 parent 5887279
commit 7facff5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/wenet/transformer/asr_model.py b/wenet/transformer/asr_model.py
@@ -54,9 +54,9 @@ def __init__(
         super().__init__()
         # note that eos is the same as sos (equivalent ID)
         self.sos = (vocab_size - 1 if special_tokens is None else
-                    special_tokens.get("sos", vocab_size - 1))
+                    special_tokens.get("<sos>", vocab_size - 1))
         self.eos = (vocab_size - 1 if special_tokens is None else
-                    special_tokens.get("eos", vocab_size - 1))
+                    special_tokens.get("<eos>", vocab_size - 1))
         self.vocab_size = vocab_size
         self.special_tokens = special_tokens
         self.ignore_id = ignore_id