Update README.md

README.md

@@ -26,6 +26,7 @@ this model can also run on 4 GB GPU RAM and know dialogs as well
 
 ```python
 
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from IPython.display import clear_output
 import textwrap
@@ -35,7 +36,6 @@ 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')])
 
 
@@ -43,25 +43,27 @@ def ppp(text:str):
   """
   pre processing prompt
   """
-  return f"<|prompter|>{text}<|endoftext|><|assistant|>"
+  return f"<|prompter|> {text} <|endoftext|><|assistant|>"
 
-def generate(text,max_new_tokens:int=512,use_ppp:bool=False,b_pair=False):
+def generate(text,max_new_tokens:int=1024,use_ppp:bool=False,b_pair=False):
   text = ppp(text) if use_ppp else text
   
   for i in range(max_new_tokens):
-    enc = tokenizer(text,return_tensors='pt')
+    enc = tokenizer(text,return_tensors='pt',add_special_tokens=False)
     text_r = text
-    enc = model.generate(**enc,max_new_tokens=1,pad_token_id=0)
-    text = tokenizer.decode(enc[0])
-    if text.endswith(tokenizer.eos_token):
+    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'):
       break
     else:
       yield text[len(text_r):] if b_pair else text
 
-for v in generate('where is empire building ?',512,True):
+
+for v in generate('what is a gpu',512,True):
   clear_output(wait=True)
   print(verify_text(v),end='')
 
+
 ```
 
 # Pythia-1B