Update app.py

app.py

@@ -21,16 +21,17 @@ def format_prompt(message, history):
   return prompt
 
 def generate(
-    prompt, history, system_prompt="You are Ultron from the Marvel Cinematic Universe and you have all his characteristics", temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.5,
+    prompt, history, system_prompt="You are Ultron from the Marvel Cinematic Universe and you have all his characteristics", temperature=0.9, max_length=50, top_p=0.95, repetition_penalty=1.5,
 ):
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
     top_p = float(top_p)
 
+    # Generate a longer response initially
     generate_kwargs = dict(
         temperature=temperature,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=max_length * 2, # Double the max_length to generate a longer response
         top_p=top_p,
         repetition_penalty=repetition_penalty,
         do_sample=True,
@@ -43,6 +44,10 @@ def generate(
 
     for response in stream:
         output += response.token.text
+        # Truncate the output to the desired max_length
+        if len(output) > max_length:
+            output = output[:max_length].rsplit(' ', 1)[0] + '...'
+            break
         yield output
     return output