Updated README

README.md

@@ -43,6 +43,7 @@ git clone https://github.com/mit-han-lab/llm-awq \
 ```
 
 ```python
+import time
 import torch
 from awq.quantize.quantizer import real_quantize_model_weight
 from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer, TextStreamer
@@ -84,6 +85,7 @@ prompt = f'''What is the difference between nuclear fusion and fission?
 ###Response:'''
 
 input_ids = tokenizer(prompt, return_tensors='pt').input_ids.cuda()
+t1 = time.time()
 output = model.generate(
     inputs=input_ids, 
     temperature=0.7,
@@ -93,6 +95,9 @@ output = model.generate(
     repetition_penalty=1.1,
     eos_token_id=tokenizer.eos_token_id,
     streamer=streamer)
+t2 = time.time()
+print("*"*80)
+print(f"Generated {num_tokens/(t2-t1):.2f} token/s; {(t2-t1)*1000/num_tokens:.2f} ms/token")
 ```
 
 ## Evaluation