Update app.py

app.py

@@ -13,28 +13,23 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 n_gpu = 0 if torch.cuda.is_available()==False else torch.cuda.device_count()
 model.to(device)
 
-early_stop_pattern = tok.eos_token
-print(f'Early stop pattern = \"{early_stop_pattern}\"')
-
 def generate(text = ""):  
   streamer = TextIteratorStreamer(tok, timeout=10.)
   if len(text) == 0:
     text = " "
   inputs = tok([text], return_tensors="pt").to(device)
-  generation_kwargs = dict(inputs, streamer=streamer, repetition_penalty=2.0, do_sample=True, top_k=40, top_p=0.97, max_new_tokens=128)
+  generation_kwargs = dict(inputs, streamer=streamer, repetition_penalty=2.0, do_sample=True, top_k=40, top_p=0.97, max_new_tokens=128, pad_token_id = model.config.eos_token_id, early_stopping=True, no_repeat_ngram_size=4)
   thread = Thread(target=model.generate, kwargs=generation_kwargs)
   thread.start()
   generated_text = ""
   for new_text in streamer:
-    yield generated_text + new_text
-    #print(new_text, end ="")
+    yield generated_text + new_text    
     generated_text += new_text
     if early_stop_pattern in generated_text:
-      generated_text = generated_text[: generated_text.find(early_stop_pattern) if early_stop_pattern else None]
+      generated_text = generated_text[: generated_text.find(tok.eos_token) if tok.eos_token else None]
       streamer.end()
-      #print("\n--\n")
       yield generated_text
-      return
+  return generated_text
 
 demo = gr.Interface(
     title="TextIteratorStreamer + Gradio demo",