webis-de#6 tokens als liste

lisalehna · Dec 5, 2024 · b911165 · b911165
1 parent 6c56367
commit b911165
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/lightning_ir/models/mvr/tokenizer.py b/lightning_ir/models/mvr/tokenizer.py
@@ -7,11 +7,11 @@ def __init__(
     ):
         super().__init__()
         if num_viewer_tokens is not None:
-            viewer_tokens = " ".join(f"[VIE{idx}]" for idx in range(num_viewer_tokens))
+            viewer_tokens = [f"[VIE{idx}]" for idx in range(num_viewer_tokens)]
             self.add_tokens(viewer_tokens, special_tokens=True)
-            viewer_token_ids = [(f"[VIE{viewer_token_id}]", self.viewer_token_id(f"[VIE{viewer_token_id}]")) for viewer_token_id in range(num_viewer_tokens)]
+            viewer_token_ids = [(viewer_tokens[viewer_token_id], self.viewer_token_id(viewer_token_id)) for viewer_token_id in range(num_viewer_tokens)]
             self.doc_post_processor = TemplateProcessing(
-                single=f"{viewer_tokens} $0 [SEP]",
+                single=f"{" ".join(viewer_tokens)} $0 [SEP]",
                 pair=f"[CLS] {self.QUERY_TOKEN} $A [SEP] $B:1 [SEP]:1",
                 special_tokens=[
                     ("[CLS]", self.cls_token_id),