Update Assistant input handling

artitw · Jan 14, 2024 · 7e1ac66 · 7e1ac66
1 parent 1c78d23
commit 7e1ac66
Show file tree

Hide file tree

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

diff --git a/text2text/assistant.py b/text2text/assistant.py
@@ -36,6 +36,8 @@ def completion_tokens(self, input_lines):
     return [len(x) for x in input_ids]
 
   def transform(self, input_lines, retriever=None, **kwargs):
+    if isinstance(input_lines, str):
+      input_lines = [input_lines]
     df = self.completion_preprocess(input_lines, retriever, **kwargs)
     temperature = kwargs.get('temperature', 0.7)
     top_p = kwargs.get('top_p', 0.95)