File size: 2,739 Bytes
ad68b7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """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()
|