| import torch, time | |
| def tokens_per_sec(model, tok): | |
| inp = tok("Hello world", return_tensors="pt") | |
| start=time.time() | |
| out = model.generate(**inp, max_new_tokens=16) | |
| elapsed=max(time.time()-start,1e-6) | |
| toks = out.shape[-1]-inp['input_ids'].shape[-1] | |
| return toks/elapsed | |