code/chat.py

import argparse
import os

import torch
from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast

from config import ViuAIConfig
from model import ViuAI


def strip_training_prefixes(state_dict):
    """Make checkpoints saved with DDP and/or torch.compile loadable."""
    cleaned = {}
    for key, value in state_dict.items():
        for prefix in ("module._orig_mod.", "_orig_mod.", "module."):
            if key.startswith(prefix):
                key = key[len(prefix):]
                break
        cleaned[key] = value
    return cleaned


def load_tokenizer(repo_id, token):
    tokenizer_path = hf_hub_download(
        repo_id=repo_id,
        filename="tokenizer/tokenizer.json",
        token=token,
    )
    return PreTrainedTokenizerFast(
        tokenizer_file=tokenizer_path,
        pad_token="<pad>",
        bos_token="<bos>",
        eos_token="<eos>",
        unk_token="<unk>",
        additional_special_tokens=["<|user|>", "<|assistant|>", "<|endofturn|>", "<think>", "</think>"],
    )


def download_checkpoint(repo_id, checkpoint_kind, sft_checkpoint_dir, token):
    candidates = (
        [
            f"{sft_checkpoint_dir}/sft_ckpt_best.pt",
            f"{sft_checkpoint_dir}/sft_ckpt_latest.pt",
            # Backward-compatible fallback for the original flat layout.
            "sft_checkpoints/sft_ckpt_best.pt",
            "sft_checkpoints/sft_ckpt_latest.pt",
        ]
        if checkpoint_kind == "sft"
        else ["checkpoints/ckpt_latest.pt", "checkpoints/ckpt_best.pt"]
    )
    errors = []
    for filename in candidates:
        try:
            return hf_hub_download(repo_id=repo_id, filename=filename, token=token), filename
        except Exception as exc:
            errors.append(f"{filename}: {type(exc).__name__}")
    raise RuntimeError("Could not download a checkpoint. Tried " + "; ".join(errors))


def main():
    parser = argparse.ArgumentParser(description="Chat with ViuAI-500M.")
    parser.add_argument("--repo-id", default=os.environ.get("VIUAI_MODEL_REPO", "ViuAI/ViuAI-500M"))
    parser.add_argument("--checkpoint", choices=("sft", "base"), default="sft")
    parser.add_argument(
        "--sft-checkpoint-dir",
        default=os.environ.get("SFT_CHECKPOINT_DIR", "sft_checkpoints/sft_v1"),
        help="Hugging Face directory containing the SFT checkpoints.",
    )
    args = parser.parse_args()

    token = os.environ.get("HF_TOKEN")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Loading ViuAI-500M on {device.upper()} ({args.checkpoint} checkpoint)...")

    tokenizer = load_tokenizer(args.repo_id, token)
    checkpoint_path, checkpoint_name = download_checkpoint(
        args.repo_id, args.checkpoint, args.sft_checkpoint_dir.strip("/"), token
    )
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

    vocab_size = 64005 if args.checkpoint == "sft" else 64000
    model = ViuAI(ViuAIConfig(vocab_size=vocab_size, use_checkpoint=False)).to(device)
    model.load_state_dict(strip_training_prefixes(checkpoint["model"]), strict=True)
    model.head.weight = model.tok_emb.weight
    model.eval()
    print(f"Ready: {checkpoint_name}. Type 'quit' or 'exit' to stop.")

    end_of_turn_id = tokenizer.convert_tokens_to_ids("<|endofturn|>")
    while True:
        try:
            message = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if message.lower() in {"quit", "exit"}:
            break
        if not message:
            continue

        # This must match the special-token format used while building SFT data.
        prompt = f"<|user|>{message}<|endofturn|><|assistant|>"
        input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
        output_ids = model.generate(
            input_ids,
            max_new_tokens=256,
            temperature=0.7,
            top_k=40,
            top_p=0.9,
            repetition_penalty=1.1,
            eos_token_id=end_of_turn_id,
        )
        # Decode with special tokens to catch <think> tags
        reply = tokenizer.decode(output_ids[0, input_ids.size(1):], skip_special_tokens=False)
        reply = reply.replace("<|endofturn|>", "")
        
        import re
        think_match = re.search(r'<think>(.*?)</think>', reply, re.DOTALL)
        if think_match:
            thinking = think_match.group(1).strip()
            reply = re.sub(r'<think>.*?</think>', '', reply, flags=re.DOTALL).strip()
            print(f"\033[90m🤔 [Thinking...]\n{thinking}\033[0m\n" + "-"*40)
            
        print(f"🤖 ViuAI: {reply.strip()}\n")


if __name__ == "__main__":
    main()