"""Give the same prompt to a base checkpoint and to its instruction-tuned sibling. Purpose: show that a chat assistant is the same next-token loop with a template around it, by running one question through a base model as raw text and through the instruction-tuned model through its own chat template, printing the template's control tokens, the tokens it adds, and how many generated tokens went into the thinking block and the answer. Platform: all (cuda on Tracks S and N, ROCm on Track X, mps on Track M, or the CPU; the device is chosen automatically and printed) Minimum memory: 8 GB (one checkpoint of about 3.4 GB in BF16 is resident at a time) Assumes: both checkpoints downloaded with `hf download ... --local-dir `, and torch, transformers and accelerate installed in the course environment from Part 1. The two models are loaded one at a time and released, so only one is resident. Usage: python base-vs-instruct.py --base ~/llm-course/models/qwen3-1.7b-base --instruct ~/llm-course/models/qwen3-1.7b [--prompt "..."] [--max-new-tokens 80] [--no-thinking] [--labbook labbook.md] """ import argparse import gc import json import time from pathlib import Path import torch from transformers import AutoModelForCausalLM, AutoTokenizer THINK_START = "" THINK_END = "" def pick_device() -> torch.device: if torch.cuda.is_available(): return torch.device("cuda") mps = getattr(torch.backends, "mps", None) if mps is not None and mps.is_available(): return torch.device("mps") return torch.device("cpu") def release(model) -> None: """Drop a model and give the memory back before loading the next one.""" del model gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def check_dir(path: str, label: str) -> None: p = Path(path).expanduser() if p.is_dir() and not (p / "config.json").exists(): raise SystemExit(f"--{label} {p} has no config.json; the download did not finish, or the path is wrong") def run(path: str, prompt_ids, device: torch.device, dtype, max_new_tokens: int, tokeniser) -> dict: """Load a checkpoint, generate greedily from prepared ids, print and release.""" model = AutoModelForCausalLM.from_pretrained(path, dtype=dtype).to(device) model.eval() started = time.perf_counter() with torch.no_grad(): out = model.generate( **prompt_ids.to(device), do_sample=False, max_new_tokens=max_new_tokens, pad_token_id=tokeniser.pad_token_id or tokeniser.eos_token_id, ) seconds = time.perf_counter() - started prompt_len = prompt_ids["input_ids"].shape[1] generated = out[0, prompt_len:].tolist() stopped_on_eos = len(generated) < max_new_tokens or generated[-1] in {tokeniser.eos_token_id, tokeniser.pad_token_id} release(model) return {"ids": generated, "seconds": seconds, "stopped_on_eos": stopped_on_eos, "text": tokeniser.decode(generated, skip_special_tokens=True)} def split_thinking(generated: list[int], tokeniser) -> tuple[str, str, int, int, bool]: """Split a Qwen3 reply at the token, as the model card describes. Returns (thinking, answer, thinking_tokens, answer_tokens, closed). When the reply opened a block that never closed, the whole reply is thinking and closed is False: the budget ran out inside the block. """ think_start_id = tokeniser.convert_tokens_to_ids(THINK_START) think_end_id = tokeniser.convert_tokens_to_ids(THINK_END) if think_end_id is not None and think_end_id in generated: cut = len(generated) - generated[::-1].index(think_end_id) thinking = tokeniser.decode(generated[:cut], skip_special_tokens=True).strip("\n") answer = tokeniser.decode(generated[cut:], skip_special_tokens=True).strip("\n") return thinking, answer, cut, len(generated) - cut, True if think_start_id is not None and think_start_id in generated: thinking = tokeniser.decode(generated, skip_special_tokens=True).strip("\n") return thinking, "", len(generated), 0, False return "", tokeniser.decode(generated, skip_special_tokens=True).strip("\n"), 0, len(generated), True def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--base", required=True, help="base checkpoint directory or Hub id") parser.add_argument("--instruct", required=True, help="instruction-tuned checkpoint directory or Hub id") parser.add_argument("--prompt", default="What is the capital of Brazil?") parser.add_argument("--max-new-tokens", type=int, default=80) parser.add_argument("--no-thinking", action="store_true", help="pass enable_thinking=False to the chat template (a Qwen3 template option)") parser.add_argument("--dtype", default="auto", choices=["auto", "float32", "bfloat16", "float16"]) parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file") args = parser.parse_args() check_dir(args.base, "base") check_dir(args.instruct, "instruct") device = pick_device() # bfloat16 is how these checkpoints were saved; --dtype float32 needs twice the memory. dtype = torch.bfloat16 if args.dtype == "auto" else getattr(torch, args.dtype) print(f"device: {device} dtype: {dtype} torch {torch.__version__} greedy decoding (do_sample=False)") # ---- the base model sees the prompt as ordinary text ----------------------------- base_tokeniser = AutoTokenizer.from_pretrained(args.base) base_inputs = base_tokeniser(args.prompt, return_tensors="pt") base_len = int(base_inputs["input_ids"].shape[1]) print(f"\n=== base: {args.base}") print(f"input as sent : {args.prompt!r}") print(f"input tokens : {base_len} {base_tokeniser.convert_ids_to_tokens(base_inputs['input_ids'][0].tolist())}") base = run(args.base, base_inputs, device, dtype, args.max_new_tokens, base_tokeniser) print(f"generated : {len(base['ids'])} tokens in {base['seconds']:.1f} s; stopped on end-of-text: {base['stopped_on_eos']}") print(f"continuation : {base['text']!r}") # ---- the instruction-tuned model gets the same question through its template ------ chat_tokeniser = AutoTokenizer.from_pretrained(args.instruct) messages = [{"role": "user", "content": args.prompt}] template_kwargs = {"add_generation_prompt": True} if args.no_thinking: template_kwargs["enable_thinking"] = False rendered = chat_tokeniser.apply_chat_template(messages, tokenize=False, **template_kwargs) chat_inputs = chat_tokeniser.apply_chat_template( messages, tokenize=True, return_dict=True, return_tensors="pt", **template_kwargs, ) chat_ids = chat_inputs["input_ids"][0].tolist() print(f"\n=== instruct: {args.instruct}") print(f"input as sent : the template's own string, control tokens included" f"{' (enable_thinking=False)' if args.no_thinking else ''}") print("---") print(rendered, end="") print("---") print(f"input tokens : {len(chat_ids)} {chat_tokeniser.convert_ids_to_tokens(chat_ids)}") print(f"template adds : {len(chat_ids) - base_len} tokens around the {base_len} of the question") chat = run(args.instruct, chat_inputs, device, dtype, args.max_new_tokens, chat_tokeniser) thinking, answer, n_think, n_answer, closed = split_thinking(chat["ids"], chat_tokeniser) print(f"generated : {len(chat['ids'])} tokens in {chat['seconds']:.1f} s; stopped on end-of-turn: {chat['stopped_on_eos']}") print(f"thinking block : {n_think} tokens{'' if closed else ' (never closed: the budget ran out inside it)'} {thinking!r}") print(f"answer : {n_answer} tokens {answer!r}") if args.labbook: record = { "lab": "part-02/base-vs-instruct", "base": args.base, "instruct": args.instruct, "device": str(device), "dtype": str(dtype), "torch": torch.__version__, "prompt": args.prompt, "max_new_tokens": args.max_new_tokens, "enable_thinking": not args.no_thinking, "base_input_tokens": base_len, "chat_input_tokens": len(chat_ids), "template_overhead_tokens": len(chat_ids) - base_len, "base_generated_tokens": len(base["ids"]), "base_seconds": round(base["seconds"], 1), "base_completion": base["text"], "chat_generated_tokens": len(chat["ids"]), "chat_seconds": round(chat["seconds"], 1), "chat_thinking_tokens": n_think, "chat_thinking_closed": closed, "chat_answer_tokens": n_answer, "chat_thinking": thinking, "chat_answer": answer, } with Path(args.labbook).expanduser().open("a", encoding="utf-8") as fh: fh.write(json.dumps(record, ensure_ascii=False) + "\n") print(f"\nrecorded in {args.labbook}") if __name__ == "__main__": main()