Merge pull request #302 from HazyResearch/flash-fill-max-tokens

add max tokens option to flash fill
HazyResearch · Mar 12, 2023 · f6fe43f · f6fe43f
2 parents 4ee7f36 + f0b30f9
commit f6fe43f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/meerkat/interactive/app/src/lib/component/contrib/flash_fill/__init__.py b/meerkat/interactive/app/src/lib/component/contrib/flash_fill/__init__.py
@@ -23,7 +23,10 @@ def __init__(
         df: "DataFrame",
         target_column: str,
         manifest_cache_dir: str = "~/.cache/manifest",
+        max_tokens: int = 1,
     ):
+        self.max_tokens = max_tokens
+
         df = df.view()
         if target_column not in df.columns:
             df[target_column] = ""
@@ -139,7 +142,7 @@ def run_manifest(
                 client_connection=os.getenv("OPENAI_API_KEY"),
                 engine=engine,
                 temperature=0,
-                max_tokens=1,
+                max_tokens=self.max_tokens,
                 cache_name="sqlite",
                 cache_connection=os.path.join(self.manifest_cache_dir, "cache.sqlite"),
             )