allow optional schema_return

dbanys · May 7, 2021 · 0c0a099 · 0c0a099
1 parent c648e20
commit 0c0a099
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/aitextgen/aitextgen.py b/aitextgen/aitextgen.py
@@ -363,7 +363,7 @@ def generate(
             # Schema token handling
             if schema:
                 schema_tokens = getattr(self.model.config, "schema_tokens")
-                schema_return = getattr(self.model.config, "schema_return")
+                schema_return = getattr(self.model.config, "schema_return", None)
                 schema_tokens_enc = self.tokenizer(text=schema_tokens)["input_ids"]
 
                 nonalphanum_pattern = re.compile(r"[\W_]+", re.UNICODE)