初始化项目，由ModelHub XC社区提供模型
Model: divakar-yadav/transformer-1b-chat Source: Original Platform
This commit is contained in:
ModelHub XC
318
training_code/chat.py
Normal file
318
training_code/chat.py
Normal file
@@ -0,0 +1,318 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Interactive chat with the 1B Transformer.
|
||||
Runs in an infinite conversation loop from the terminal.
|
||||
|
||||
Usage:
|
||||
python chat.py # auto-find latest checkpoint
|
||||
python chat.py /jfs/deepak-kumar/checkpoints/step_19000.pt # specific checkpoint
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import glob
|
||||
import time
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import readline # enables arrow keys and history in input()
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from model.config import ModelConfig
|
||||
from model.transformer import Transformer
|
||||
from model.data import get_tokenizer
|
||||
|
||||
|
||||
def find_latest_checkpoint():
|
||||
"""Look for DPO > SFT > pretrained checkpoint."""
|
||||
dpo_dir = "/jfs/deepak-kumar/checkpoints_dpo"
|
||||
sft_dir = "/jfs/deepak-kumar/checkpoints_sft"
|
||||
pt_dir = "/jfs/deepak-kumar/checkpoints"
|
||||
|
||||
# Prefer DPO final
|
||||
dpo_final = os.path.join(dpo_dir, "dpo_final.pt")
|
||||
if os.path.exists(dpo_final):
|
||||
return dpo_final, True
|
||||
|
||||
dpo_files = glob.glob(os.path.join(dpo_dir, "dpo_step_*.pt"))
|
||||
if dpo_files:
|
||||
return max(dpo_files, key=lambda f: int(f.split("dpo_step_")[1].split(".")[0])), True
|
||||
|
||||
# Then SFT
|
||||
sft_final = os.path.join(sft_dir, "sft_final.pt")
|
||||
if os.path.exists(sft_final):
|
||||
return sft_final, True
|
||||
|
||||
sft_files = glob.glob(os.path.join(sft_dir, "sft_step_*.pt"))
|
||||
if sft_files:
|
||||
return max(sft_files, key=lambda f: int(f.split("sft_step_")[1].split(".")[0])), True
|
||||
|
||||
# Fall back to pretrained
|
||||
pt_files = glob.glob(os.path.join(pt_dir, "step_*.pt"))
|
||||
if pt_files:
|
||||
return max(pt_files, key=lambda f: int(os.path.basename(f).split("_")[1].split(".")[0])), False
|
||||
|
||||
return None, False
|
||||
|
||||
|
||||
def load_model(checkpoint_path, tokenizer, device="cuda:0"):
|
||||
config = ModelConfig()
|
||||
model = Transformer(config)
|
||||
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||||
|
||||
# Handle expanded vocab from SFT
|
||||
saved_vocab = ckpt.get("vocab_size", config.vocab_size)
|
||||
if saved_vocab > config.vocab_size:
|
||||
config.vocab_size = saved_vocab
|
||||
model = Transformer(config)
|
||||
|
||||
model.load_state_dict(ckpt["model"])
|
||||
model = model.to(device).bfloat16().eval()
|
||||
step = ckpt.get("step", "?")
|
||||
loss = ckpt.get("loss", "?")
|
||||
del ckpt
|
||||
torch.cuda.empty_cache()
|
||||
return model, config, step, loss
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_stream(model, tokenizer, prompt, max_new_tokens=512,
|
||||
temperature=0.8, top_k=50, top_p=0.9,
|
||||
repetition_penalty=1.15, device="cuda:0",
|
||||
stop_token_ids=None):
|
||||
"""Generate tokens one at a time, yielding each for streaming output."""
|
||||
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
|
||||
generated_ids = []
|
||||
prev_decoded_len = 0
|
||||
|
||||
if stop_token_ids is None:
|
||||
stop_token_ids = set()
|
||||
else:
|
||||
stop_token_ids = set(stop_token_ids)
|
||||
stop_token_ids.add(tokenizer.eos_token_id)
|
||||
|
||||
for _ in range(max_new_tokens):
|
||||
if input_ids.shape[1] >= model.config.max_seq_len:
|
||||
break
|
||||
|
||||
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
||||
logits, _ = model(input_ids)
|
||||
|
||||
logits = logits[:, -1, :]
|
||||
|
||||
if repetition_penalty != 1.0 and generated_ids:
|
||||
prev_tokens = torch.tensor(generated_ids, device=device).unique()
|
||||
for token_id in prev_tokens:
|
||||
if logits[0, token_id] > 0:
|
||||
logits[0, token_id] /= repetition_penalty
|
||||
else:
|
||||
logits[0, token_id] *= repetition_penalty
|
||||
|
||||
logits = logits / temperature
|
||||
|
||||
if top_k > 0:
|
||||
topk_vals, _ = torch.topk(logits, top_k)
|
||||
logits[logits < topk_vals[:, -1:]] = float("-inf")
|
||||
|
||||
if top_p < 1.0:
|
||||
sorted_logits, sorted_idx = torch.sort(logits, descending=True)
|
||||
cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
||||
mask = cum_probs - F.softmax(sorted_logits, dim=-1) >= top_p
|
||||
sorted_logits[mask] = float("-inf")
|
||||
logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
|
||||
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
next_token = torch.multinomial(probs, num_samples=1)
|
||||
token_id = next_token.item()
|
||||
|
||||
# Stop on any stop token (EOS, <|end|>, <|user|>)
|
||||
if token_id in stop_token_ids:
|
||||
break
|
||||
|
||||
generated_ids.append(token_id)
|
||||
input_ids = torch.cat([input_ids, next_token], dim=1)
|
||||
|
||||
full_decoded = tokenizer.decode(generated_ids, skip_special_tokens=True)
|
||||
new_text = full_decoded[prev_decoded_len:]
|
||||
prev_decoded_len = len(full_decoded)
|
||||
yield new_text
|
||||
|
||||
return
|
||||
|
||||
|
||||
def print_banner(step, loss, device):
|
||||
print("\033[1;36m") # cyan bold
|
||||
print("=" * 60)
|
||||
print(" 1B TRANSFORMER — Interactive Chat")
|
||||
print("=" * 60)
|
||||
print(f"\033[0m Checkpoint : step {step}")
|
||||
print(f" Loss : {loss}")
|
||||
print(f" Device : {device}")
|
||||
print(f" Parameters : 1.106B")
|
||||
print()
|
||||
print(" \033[90mCommands:\033[0m")
|
||||
print(" \033[33m/quit\033[0m — exit")
|
||||
print(" \033[33m/clear\033[0m — clear conversation context")
|
||||
print(" \033[33m/temp N\033[0m — set temperature (default 0.8)")
|
||||
print(" \033[33m/tokens N\033[0m — set max tokens (default 512)")
|
||||
print(" \033[33m/topp N\033[0m — set top-p (default 0.9)")
|
||||
print(" \033[33m/topk N\033[0m — set top-k (default 50)")
|
||||
print(" \033[33m/rep N\033[0m — set repetition penalty (default 1.15)")
|
||||
print()
|
||||
print("\033[90m" + "─" * 60 + "\033[0m")
|
||||
|
||||
|
||||
def main():
|
||||
device = "cuda:0"
|
||||
|
||||
is_sft = False
|
||||
if len(sys.argv) > 1:
|
||||
checkpoint = sys.argv[1]
|
||||
is_sft = "sft" in checkpoint.lower()
|
||||
else:
|
||||
result = find_latest_checkpoint()
|
||||
if result[0] is None:
|
||||
print("No checkpoint found!")
|
||||
sys.exit(1)
|
||||
checkpoint, is_sft = result
|
||||
|
||||
tokenizer = get_tokenizer()
|
||||
|
||||
# Add chat tokens for SFT models
|
||||
if is_sft:
|
||||
special_tokens = ["<|user|>", "<|assistant|>", "<|end|>"]
|
||||
vocab = tokenizer.get_vocab()
|
||||
new_tokens = [t for t in special_tokens if t not in vocab]
|
||||
if new_tokens:
|
||||
tokenizer.add_tokens(new_tokens, special_tokens=True)
|
||||
|
||||
print(f"\n Loading model from {checkpoint}...")
|
||||
print(f" Mode: {'SFT (chat)' if is_sft else 'Base (completion)'}")
|
||||
model, config, step, loss = load_model(checkpoint, tokenizer, device)
|
||||
print(f" Model loaded!\n")
|
||||
|
||||
print_banner(step, loss, device)
|
||||
if is_sft:
|
||||
print(" \033[1;32mSFT mode: The model will respond as a chat assistant.\033[0m\n")
|
||||
|
||||
# Settings
|
||||
temperature = 0.7 if is_sft else 0.8
|
||||
max_tokens = 512
|
||||
top_p = 0.9
|
||||
top_k = 50
|
||||
rep_penalty = 1.15
|
||||
context = ""
|
||||
|
||||
# Chat template tokens for SFT
|
||||
USER_START = "<|user|>\n"
|
||||
ASST_START = "<|assistant|>\n"
|
||||
TURN_END = "\n<|end|>\n"
|
||||
|
||||
# Build stop token IDs for generation
|
||||
sft_stop_ids = []
|
||||
if is_sft:
|
||||
vocab = tokenizer.get_vocab()
|
||||
for tok_str in ["<|end|>", "<|user|>"]:
|
||||
if tok_str in vocab:
|
||||
sft_stop_ids.append(vocab[tok_str])
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\n\033[1;32mYou:\033[0m ").strip()
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\n\nGoodbye!")
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
# Handle commands
|
||||
if user_input.startswith("/"):
|
||||
cmd = user_input.lower().split()
|
||||
if cmd[0] == "/quit":
|
||||
print("Goodbye!")
|
||||
break
|
||||
elif cmd[0] == "/clear":
|
||||
context = ""
|
||||
print("\033[90m [Context cleared]\033[0m")
|
||||
continue
|
||||
elif cmd[0] == "/temp" and len(cmd) > 1:
|
||||
temperature = float(cmd[1])
|
||||
print(f"\033[90m [Temperature set to {temperature}]\033[0m")
|
||||
continue
|
||||
elif cmd[0] == "/tokens" and len(cmd) > 1:
|
||||
max_tokens = int(cmd[1])
|
||||
print(f"\033[90m [Max tokens set to {max_tokens}]\033[0m")
|
||||
continue
|
||||
elif cmd[0] == "/topp" and len(cmd) > 1:
|
||||
top_p = float(cmd[1])
|
||||
print(f"\033[90m [Top-p set to {top_p}]\033[0m")
|
||||
continue
|
||||
elif cmd[0] == "/topk" and len(cmd) > 1:
|
||||
top_k = int(cmd[1])
|
||||
print(f"\033[90m [Top-k set to {top_k}]\033[0m")
|
||||
continue
|
||||
elif cmd[0] == "/rep" and len(cmd) > 1:
|
||||
rep_penalty = float(cmd[1])
|
||||
print(f"\033[90m [Repetition penalty set to {rep_penalty}]\033[0m")
|
||||
continue
|
||||
else:
|
||||
print("\033[90m Unknown command. Try /quit, /clear, /temp, /tokens, /topp, /topk, /rep\033[0m")
|
||||
continue
|
||||
|
||||
# Build prompt
|
||||
if is_sft:
|
||||
prompt = context + USER_START + user_input + TURN_END + ASST_START
|
||||
else:
|
||||
if context:
|
||||
prompt = context + "\n" + user_input
|
||||
else:
|
||||
prompt = user_input
|
||||
|
||||
# Trim context if too long
|
||||
while len(tokenizer.encode(prompt)) > config.max_seq_len - max_tokens:
|
||||
if is_sft:
|
||||
parts = context.split(TURN_END)
|
||||
if len(parts) <= 2:
|
||||
break
|
||||
context = TURN_END.join(parts[2:])
|
||||
prompt = context + USER_START + user_input + TURN_END + ASST_START
|
||||
else:
|
||||
lines = prompt.split("\n")
|
||||
if len(lines) <= 2:
|
||||
break
|
||||
prompt = "\n".join(lines[1:])
|
||||
|
||||
# Generate with streaming
|
||||
print("\033[1;34mModel:\033[0m ", end="", flush=True)
|
||||
t0 = time.time()
|
||||
full_response = ""
|
||||
token_count = 0
|
||||
|
||||
for token_text in generate_stream(
|
||||
model, tokenizer, prompt,
|
||||
max_new_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
repetition_penalty=rep_penalty,
|
||||
device=device,
|
||||
stop_token_ids=sft_stop_ids if is_sft else None,
|
||||
):
|
||||
print(token_text, end="", flush=True)
|
||||
full_response += token_text
|
||||
token_count += 1
|
||||
|
||||
elapsed = time.time() - t0
|
||||
tps = token_count / max(elapsed, 1e-9)
|
||||
print(f"\n\033[90m [{token_count} tokens, {tps:.1f} tok/s, {elapsed:.1f}s]\033[0m")
|
||||
|
||||
# Append to context for multi-turn
|
||||
if is_sft:
|
||||
context = (context + USER_START + user_input + TURN_END +
|
||||
ASST_START + full_response.strip() + TURN_END)
|
||||
else:
|
||||
context = prompt + full_response
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user