Update README.md

README.md

@@ -36,7 +36,8 @@ tokenizer = AutoTokenizer.from_pretrained("erfanzar/PGT-1B-2EP")
 
 model = AutoModelForCausalLM.from_pretrained("erfanzar/PGT-1B-2EP",device_map='auto',load_in_8bit=True)
 
-verify_text = lambda txt : '\n'.join([textwrap.fill(txt, width=140) for txt in txt.split('\n')])
+
+verify_text = lambda txt : '\n'.join([textwrap.fill(txt, width=110) for txt in txt.split('\n')])
 
 
 def ppp(text:str):
@@ -52,8 +53,11 @@ def generate(text,max_new_tokens:int=1024,use_ppp:bool=False,b_pair=False):
     enc = tokenizer(text,return_tensors='pt',add_special_tokens=False)
     text_r = text
     enc = model.generate(enc.input_ids,max_new_tokens=1,pad_token_id=0)
-    text = tokenizer.decode(enc[0],skip_special_tokens=False).replace('\n\n\n\n',tokenizer.eos_token)
-    if text.endswith(tokenizer.eos_token) or text.endswith('\n\n\n\n'):
+    text = tokenizer.decode(enc[0],skip_special_tokens=False)
+    text = text[:-4]+tokenizer.eos_token if text[-4:] == '\n\n\n\n' else text
+  
+    if text.endswith(tokenizer.eos_token) or text.endswith('\n\n\n\n\n'):
+      yield text[len(text_r):] if b_pair else text
       break
     else:
       yield text[len(text_r):] if b_pair else text