Update README.md

README.md

@@ -77,16 +77,23 @@ import sys
 model_path = "./" # You can modify the path for storing the local model
 model =  AutoModelForCausalLM.from_pretrained(model_path)
 tokenizer = AutoTokenizer.from_pretrained(model_path)
+history = "" # save dialog history
+max_history_length = 1024
+
 print("Human:")
 line = input()
 while line:
-        inputs = 'Human: ' + line.strip() + '\n\nAssistant:'
-        input_ids = tokenizer(inputs, return_tensors="pt").input_ids
+        history += 'Human: ' + line.strip() + '\n\nAssistant:'
+        if len(history) > max_history_length:
+                history = history[-1024:]
+        input_ids = tokenizer(history, return_tensors="pt").input_ids
         outputs = model.generate(input_ids, max_new_tokens=200, do_sample = True, top_k = 30, top_p = 0.85, temperature = 0.35, repetition_penalty=1.2)
         rets = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-        print("Assistant:\n" + rets[0].strip().replace(inputs, ""))
+        print("Assistant:\n" + rets[0].strip().replace(history, ""))
+        history += "\n" + rets[0].strip().replace(history, "")
         print("\n------------------------------------------------\nHuman:")
         line = input()
+
 ```
 
 ## Limitations