[no ci] minor edit

amaiya · May 10, 2023 · 61bd40a · 61bd40a
1 parent 75e17e3
commit 61bd40a
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/ktrain/text/qa/generative_qa.py b/ktrain/text/qa/generative_qa.py
@@ -152,6 +152,7 @@ def query(
         marginal_relevance: bool = True,
         answer=None,
         key_filter: Optional[bool] = None,
+        show_token_usage=False,
         # get_callbacks: Callable[[str], AsyncCallbackHandler] = lambda x: [],
     ):
         """
@@ -169,6 +170,10 @@ def query(
                 answer=answer,
                 key_filter=key_filter,
             )
+            if not show_token_usage:
+                result.formatted_answer = result.formatted_answer.split("Tokens Used")[
+                    0
+                ]
             return result
         except RuntimeError:
             raise Exception(