Use line_by_line for default gen if present

minimaxir · May 1, 2021 · 342a361 · 342a361
1 parent 34e6cad
commit 342a361
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/aitextgen/TokenDataset.py b/aitextgen/TokenDataset.py
@@ -121,6 +121,7 @@ def __init__(
                 self.tokens = np.load(f)
             self.num_subsets = self.tokens.shape[0] - block_size
             self.block_size = block_size
+            self.line_by_line = line_by_line
             self.str_suffix = "via cache."
 
             logger.info(

diff --git a/aitextgen/aitextgen.py b/aitextgen/aitextgen.py
@@ -277,7 +277,7 @@ def generate(
         self,
         n: int = 1,
         prompt: str = "",
-        prepend_bos: bool = False,
+        prepend_bos: bool = None,
         min_length: int = None,
         max_length: int = 256,
         temperature: float = 0.7,
@@ -325,6 +325,9 @@ def generate(
             prompt_tensors["input_ids"].to(self.get_device()) if prompt else None
         )
 
+        if prepend_bos is None:
+            prepend_bos = getattr(self.model.config, "line_by_line", None)
+
         if prepend_bos:
             bos = torch.tensor([[self.tokenizer.bos_token_id]]).to(self.get_device())
             if prompt: