| 13 | |
| 14 | # Referenced from https://github.com/Lightning-AI/lit-llama/blob/main/generate/full.py |
| 15 | def main( |
| 16 | prompt: str = "Hello, my name is", |
| 17 | max_new_tokens: int = 50, |
| 18 | top_k: int = 200, |
| 19 | temperature: float = 0.8, |
| 20 | checkpoint_path: Optional[Path] = None, |
| 21 | tokenizer_path: Path = Path('tokenizer'), |
| 22 | ): |
| 23 | |
| 24 | checkpoint_path = normpath(join(getcwd(), checkpoint_path)) |
| 25 | print('checkpoint_path: ', checkpoint_path) |
| 26 | checkpoint_path = Path(checkpoint_path) |
| 27 | |
| 28 | tokenizer_path = normpath(join(getcwd(), tokenizer_path)) |
| 29 | print('tokenizer_path: ', tokenizer_path) |
| 30 | tokenizer_path = Path(tokenizer_path) |
| 31 | |
| 32 | assert checkpoint_path.is_file(), checkpoint_path |
| 33 | assert tokenizer_path.is_dir(), tokenizer_path |
| 34 | |
| 35 | config = Config.from_name('2.0-Pints-Upscaled') |
| 36 | model = GPT(config) |
| 37 | checkpoint = torch.load(checkpoint_path) |
| 38 | checkpoint = checkpoint["model"] |
| 39 | model.load_state_dict(checkpoint) |
| 40 | model.eval() |
| 41 | |
| 42 | fabric = lightning.Fabric(devices=1, precision='bf16-true') |
| 43 | # fabric = lightning.Fabric(devices=1, precision='32-true') |
| 44 | model = fabric.setup(model) |
| 45 | |
| 46 | tokenizer = Tokenizer(tokenizer_path) |
| 47 | |
| 48 | prompt = f'''<|im_start|>system |
| 49 | you are an expert in writing<|im_end|> |
| 50 | <|im_start|>user |
| 51 | {prompt}<|im_end|> |
| 52 | <|im_start|>assistant\n''' |
| 53 | |
| 54 | encoded = tokenizer.encode(prompt, bos=True, eos=False, device=fabric.device) |
| 55 | print(encoded) |
| 56 | prompt_length = encoded.size(0) |
| 57 | lightning.seed_everything(1234) |
| 58 | |
| 59 | # Use `samples` to generate a few samples. |
| 60 | # for i in range(samples): |
| 61 | t0 = time.perf_counter() |
| 62 | y = generate(model, encoded, max_new_tokens, temperature=temperature, top_k=top_k, eos_id=tokenizer.eos_id) |
| 63 | t = time.perf_counter() - t0 |
| 64 | |
| 65 | model.reset_cache() |
| 66 | |
| 67 | print(tokenizer.decode(y)) |
| 68 | tokens_generated = y.size(0) - prompt_length |
| 69 | print(f"Time for inference: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec", file=sys.stderr) |
| 70 | if fabric.device.type == "cuda": |
| 71 | print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB", file=sys.stderr) |
| 72 | |