Change max_new_tokens to max_length

app.py

@@ -30,7 +30,7 @@ model = AutoModelForCausalLM.from_pretrained(
 
 
 def generate_text(prompt: str,
-                  max_new_tokens: int = 512,
+                  max_length: int = 1024,
                   temperature: float = 0.5,
                   top_p: float = 0.95,
                   top_k: int = 50) -> str:
@@ -43,8 +43,6 @@ def generate_text(prompt: str,
     inputs = inputs.to(DEVICE)
 
     # Prepare arguments for generation
-    input_length = inputs["input_ids"].shape[-1]
-    max_new_tokens = min(max_new_tokens, WINDOW_SIZE - input_length)
     if temperature >= 1.0:
         temperature = 0.99
     elif temperature <= 0.0:
@@ -60,7 +58,7 @@ def generate_text(prompt: str,
     generation_kwargs = dict(
         inputs=inputs,
         streamer=inputs,
-        max_new_tokens=max_new_tokens,
+        max_length=max_length,
         do_sample=True,
         top_p=top_p,
         top_k=top_k,
@@ -90,7 +88,7 @@ demo = gr.Interface(
             scale=10,
         ),
         gr.Slider(
-            label="Maximum new tokens",
+            label="Maximum length of the output",
             minimum=1,
             maximum=4096,
             step=1,