Update app

app.py

@@ -30,19 +30,19 @@ model = AutoModelForCausalLM.from_pretrained(
 
 
 def generate_text(prompt: str,
-                  max_length: int = 1024,
+                  max_new_tokens: int = 512,
                   temperature: float = 0.5,
                   top_p: float = 0.95,
                   top_k: int = 50) -> str:
 
     # Encode the prompt
     inputs = tokenizer([prompt],
-                       return_tensors="pt",
-                       add_special_tokens=False,
-                       return_token_type_ids=False)
-    inputs = inputs.to(DEVICE)
+                       return_tensors='pt',
+                       add_special_tokens=False).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:
@@ -58,7 +58,7 @@ def generate_text(prompt: str,
     generation_kwargs = dict(
         inputs=inputs,
         streamer=inputs,
-        max_length=max_length,
+        max_new_tokens=max_new_tokens,
         do_sample=True,
         top_p=top_p,
         top_k=top_k,
@@ -70,9 +70,10 @@ def generate_text(prompt: str,
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    # outputs = []
-    for text in streamer:
-        return text
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+    return generated_text
 
 
 demo = gr.Interface(
@@ -88,7 +89,7 @@ demo = gr.Interface(
             scale=10,
         ),
         gr.Slider(
-            label="Maximum length of the output",
+            label="Maximum new tokens",
             minimum=1,
             maximum=4096,
             step=1,