Upload app.py

app.py

@@ -257,14 +257,14 @@ def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, blo
             
             total_step += 1
             
-        if tokenizer.eos_token_id is not None:
-            finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
+            if tokenizer.eos_token_id is not None:
+                finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
+            if finished.all():
+                break
+            generated += cur_len
+            
         if finished.all():
             break
-        generated += cur_len
-        
-    if finished.all():
-        break
         
     new_tokens = x[0, prompt_len:prompt_len + max_new_tokens].tolist()
     final_text = tokenizer.decode(new_tokens, skip_special_tokens=True)