| """Interactive chat with the finetuned model. Model loads once, then you can ask questions back-to-back.""" |
|
|
| import argparse |
| import torch |
| from pathlib import Path |
| from litgpt import Tokenizer |
| from litgpt.config import Config |
| from litgpt.model import GPT |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--checkpoint_dir", type=str, default="Base/out/finetune/custom-100m-english-instruct/final") |
| parser.add_argument("--tokenizer_dir", type=str, default="Base/checkpoints/EleutherAI/pythia-160m") |
| parser.add_argument("--max_new_tokens", type=int, default=50) |
| parser.add_argument("--temperature", type=float, default=0.2) |
| parser.add_argument("--top_k", type=int, default=1) |
| args = parser.parse_args() |
|
|
| ckpt = Path(args.checkpoint_dir) |
| print(f"Loading model from {ckpt}...") |
| cfg = Config.from_checkpoint(ckpt) |
| model = GPT(cfg) |
| sd = torch.load(str(ckpt / "lit_model.pth"), map_location="cpu", weights_only=False) |
| if "model" in sd: |
| sd = sd["model"] |
| model.load_state_dict(sd, strict=False) |
| model = model.to("cuda").eval() |
| tok = Tokenizer(Path(args.tokenizer_dir)) |
| print(f"Model loaded! ({sum(p.numel() for p in model.parameters()):,} params)") |
| print(f"Settings: temp={args.temperature}, top_k={args.top_k}, max_tokens={args.max_new_tokens}") |
| print("Type your question and press Enter. Type 'quit' or 'exit' to stop.\n") |
|
|
| while True: |
| try: |
| question = input("You: ").strip() |
| except (EOFError, KeyboardInterrupt): |
| print("\nBye!") |
| break |
| if not question: |
| continue |
| if question.lower() in ("quit", "exit", "q"): |
| print("Bye!") |
| break |
|
|
| alpaca = ( |
| "Below is an instruction that describes a task. " |
| "Write a response that appropriately completes the request.\n\n" |
| f"### Instruction:\n{question}\n\n### Response:\n" |
| ) |
| ids = tok.encode(alpaca, device="cuda").unsqueeze(0) |
| with torch.no_grad(): |
| for _ in range(args.max_new_tokens): |
| logits = model(ids[:, -cfg.block_size:]) |
| logits = logits[:, -1, :] / args.temperature |
| v, _ = torch.topk(logits, min(args.top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = float("-inf") |
| probs = torch.softmax(logits, dim=-1) |
| nxt = torch.multinomial(probs, num_samples=1) |
| ids = torch.cat([ids, nxt], dim=1) |
| if nxt.item() == tok.eos_id: |
| break |
| resp = tok.decode(ids[0]).split("### Response:")[-1].strip() |
| print(f"\nLUNA: {resp}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
