Assistant token counter

artitw · Sep 23, 2023 · dc6231b · dc6231b
1 parent f0ca8cc
commit dc6231b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
   name="text2text",
-  version="1.3.1",
+  version="1.3.2",
   author="artitw",
   author_email="artitw@gmail.com",
   description="Text2Text: Crosslingual NLP/G toolkit",

diff --git a/text2text/assistant.py b/text2text/assistant.py
@@ -37,7 +37,7 @@ def num_tokens(self, input_lines, src_lang='en'):
     df = self.preprocess(input_lines, src_lang)
     tok = self.__class__.tokenizer
     input_ids = tok(df["input_line"].tolist(), return_tensors="pt", padding=True).input_ids
-    return len(input_ids[0])
+    return [len(x) for x in input_ids]
 
   def transform(self, input_lines, src_lang='en', retriever=None, **kwargs):
     df = self.preprocess(input_lines, src_lang, retriever, **kwargs)