More gen param tweaks

minimaxir · May 17, 2020 · 1df1211 · 1df1211
1 parent 74f54a5
commit 1df1211
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/aitextgen/aitextgen.py b/aitextgen/aitextgen.py
@@ -234,7 +234,7 @@ def generate(
         self,
         n: int = 1,
         prompt: str = None,
-        max_length: int = 200,
+        max_length: int = 1024,
         temperature: float = 0.7,
         do_sample: bool = True,
         return_as_list: bool = False,

diff --git a/aitextgen/train.py b/aitextgen/train.py
@@ -191,7 +191,10 @@ def generate_sample_text(self, trainer, pl_module):
         gen_length = min(pl_module.model.config.n_positions, 256)
 
         outputs = pl_module.model.generate(
-            max_length=gen_length, do_sample=True, num_return_sequences=self.n_generate
+            max_length=gen_length,
+            do_sample=True,
+            num_return_sequences=self.n_generate,
+            temperature=0.7,
         )
         gen_texts = [
             pl_module.tokenizer.decode(output, skip_special_tokens=True)