palimpseste-max / examples /train_chat.py

Upload examples/train_chat.py with huggingface_hub

8bf3b16 verified 3 months ago

5.99 kB

#!/usr/bin/env python
"""PALIMPSESTE — Train a conversational chat model.
Trains on Q/A pairs (from corpus_chat.py or a JSON file) so the model learns
to respond to user messages. Each pair is encoded as a self-contained episode:
BOS question EOS BOS answer EOS.
Usage:
  # train on the built-in corpus
  python examples/train_chat.py --preset small --output ./chat_model
  # train on custom Q/A pairs from a JSON file: [{"q": "...", "a": "..."}, ...]
  python examples/train_chat.py --qa-file my_pairs.json --output ./chat_model
  # train at 1B-scale capacity
  python examples/train_chat.py --preset 1b --output ./chat_model_1b
  # train and push to HF Hub
  python examples/train_chat.py --preset small --output ./chat_model --push-to-hub user/palimpseste-chat
  # then chat interactively:
  python examples/chat.py --model ./chat_model
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from palimseste.lm import PalimpsesteConfig, PRESETS
from palimseste.hf import HFPalimpsesteLM
from corpus_chat import get_corpus


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Train a PALIMPSESTE conversational chat model.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("--output", "-o", type=str, required=True,
                   help="output directory for the saved model")
    p.add_argument("--preset", type=str, default="small",
                   choices=list(PRESETS.keys()),
                   help="capacity preset (default: small)")
    p.add_argument("--qa-file", type=str, default=None,
                   help="JSON file with [{q, a}, ...]; default: built-in corpus")
    p.add_argument("--D", type=int, default=None, help="override D")
    p.add_argument("--context-window", type=int, default=None,
                   help="override context window (default: 64 for chat)")
    p.add_argument("--kernel-radius", type=int, default=None,
                   help="override kernel radius")
    p.add_argument("--temperature", type=float, default=0.3,
                   help="sampling temperature (default: 0.3 for focused answers)")
    p.add_argument("--push-to-hub", type=str, default=None)
    p.add_argument("--hub-token", type=str, default=None)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--epochs", type=int, default=1,
                   help="number of passes over the corpus (default: 1; "
                        "PALIMPSESTE is append-only so >1 reinforces)")
    return p.parse_args()


def load_pairs(args: argparse.Namespace) -> list[tuple[str, str]]:
    if args.qa_file:
        with open(args.qa_file, "r", encoding="utf-8") as f:
            data = json.load(f)
        return [(d["q"], d["a"]) for d in data]
    return get_corpus()


def main() -> None:
    args = parse_args()
    pairs = load_pairs(args)

    # build config — chat benefits from a larger context window
    config = PRESETS[args.preset]
    if args.D is not None:
        config = PalimpsesteConfig(D=args.D, context_window=config.context_window,
                                   kernel_radius=config.kernel_radius)
    if args.context_window is not None:
        config.context_window = args.context_window
    elif args.context_window is None and config.context_window < 48:
        config.context_window = 64  # chat needs room for Q+A
    if args.kernel_radius is not None:
        config.kernel_radius = args.kernel_radius
    config.temperature = args.temperature

    print("=" * 70)
    print("PALIMPSESTE — conversational chat training")
    print("=" * 70)
    print(f"config: D={config.D:,}  context_window={config.context_window}  "
          f"kernel_radius={config.kernel_radius}  temperature={config.temperature}")
    print(f"Q/A pairs: {len(pairs)}")

    rng = np.random.default_rng(args.seed)
    lm = HFPalimpsesteLM(config=config, rng=rng)

    # build vocab from all questions + answers
    full_text = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full_text)
    print(f"vocab: {lm.tokenizer.vocab_size} chars")

    # train (O(1) per token, no gradient)
    import time
    total = 0
    for epoch in range(args.epochs):
        t0 = time.perf_counter()
        n = lm.train_on_qa_pairs(pairs, verbose=(args.epochs == 1))
        dt = time.perf_counter() - t0
        total += n
        print(f"  epoch {epoch+1}/{args.epochs}: {n} tokens in {dt:.2f}s "
              f"({n/max(dt,1e-9):,.0f} tok/s), |M|={len(lm.mem)}")

    print(f"\ntotal: {total} tokens, |M| = {len(lm.mem):,}")

    # test responses on a few sample questions using the Conversation layer
    print("\n--- sample responses ---")
    from palimseste.chat import Conversation
    conv = Conversation(model=lm, learn_live=True)
    conv.register_questions(pairs)
    test_inputs = [
        "bonjour",
        "qui es-tu",
        "comment tu apprends",
        "tu utilises un gpu",
        "tu utilise un gpu",  # typo
        "who are you",
        "how do you learn",
        "c quoi palimpseste",
        "merci",
    ]
    for q in test_inputs:
        resp = conv.respond(q, temperature=0.0, seed=0)
        print(f"  Q: {q}")
        print(f"  A: {resp}")
        print()
    conv.reset()  # don't save conversation history into the model

    # save
    out_dir = Path(args.output)
    print(f"saving to {out_dir} ...")
    lm.save_pretrained(out_dir)

    # push to hub
    if args.push_to_hub:
        print(f"\npushing to HF Hub: {args.push_to_hub} ...")
        url = lm.push_to_hub(args.push_to_hub, token=args.hub_token)
        print(f"pushed: {url}")

    print("\n" + "=" * 70)
    print("done. Chat with:")
    print(f"  python examples/chat.py --model {out_dir}")
    print("=" * 70)


if __name__ == "__main__":
    main()

	#!/usr/bin/env python
	"""PALIMPSESTE — Train a conversational chat model.

	Trains on Q/A pairs (from corpus_chat.py or a JSON file) so the model learns
	to respond to user messages. Each pair is encoded as a self-contained episode:
	BOS question EOS BOS answer EOS.

	Usage:
	# train on the built-in corpus
	python examples/train_chat.py --preset small --output ./chat_model

	# train on custom Q/A pairs from a JSON file: [{"q": "...", "a": "..."}, ...]
	python examples/train_chat.py --qa-file my_pairs.json --output ./chat_model

	# train at 1B-scale capacity
	python examples/train_chat.py --preset 1b --output ./chat_model_1b

	# train and push to HF Hub
	python examples/train_chat.py --preset small --output ./chat_model --push-to-hub user/palimpseste-chat

	# then chat interactively:
	python examples/chat.py --model ./chat_model
	"""

	from __future__ import annotations

	import argparse
	import json
	import sys
	from pathlib import Path

	import numpy as np

	sys.path.insert(0, str(Path(__file__).resolve().parent))
	sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

	from palimseste.lm import PalimpsesteConfig, PRESETS
	from palimseste.hf import HFPalimpsesteLM
	from corpus_chat import get_corpus


	def parse_args() -> argparse.Namespace:
	p = argparse.ArgumentParser(
	description="Train a PALIMPSESTE conversational chat model.",
	formatter_class=argparse.RawDescriptionHelpFormatter,
	epilog=__doc__,
	)
	p.add_argument("--output", "-o", type=str, required=True,
	help="output directory for the saved model")
	p.add_argument("--preset", type=str, default="small",
	choices=list(PRESETS.keys()),
	help="capacity preset (default: small)")
	p.add_argument("--qa-file", type=str, default=None,
	help="JSON file with [{q, a}, ...]; default: built-in corpus")
	p.add_argument("--D", type=int, default=None, help="override D")
	p.add_argument("--context-window", type=int, default=None,
	help="override context window (default: 64 for chat)")
	p.add_argument("--kernel-radius", type=int, default=None,
	help="override kernel radius")
	p.add_argument("--temperature", type=float, default=0.3,
	help="sampling temperature (default: 0.3 for focused answers)")
	p.add_argument("--push-to-hub", type=str, default=None)
	p.add_argument("--hub-token", type=str, default=None)
	p.add_argument("--seed", type=int, default=42)
	p.add_argument("--epochs", type=int, default=1,
	help="number of passes over the corpus (default: 1; "
	"PALIMPSESTE is append-only so >1 reinforces)")
	return p.parse_args()


	def load_pairs(args: argparse.Namespace) -> list[tuple[str, str]]:
	if args.qa_file:
	with open(args.qa_file, "r", encoding="utf-8") as f:
	data = json.load(f)
	return [(d["q"], d["a"]) for d in data]
	return get_corpus()


	def main() -> None:
	args = parse_args()
	pairs = load_pairs(args)

	# build config — chat benefits from a larger context window
	config = PRESETS[args.preset]
	if args.D is not None:
	config = PalimpsesteConfig(D=args.D, context_window=config.context_window,
	kernel_radius=config.kernel_radius)
	if args.context_window is not None:
	config.context_window = args.context_window
	elif args.context_window is None and config.context_window < 48:
	config.context_window = 64 # chat needs room for Q+A
	if args.kernel_radius is not None:
	config.kernel_radius = args.kernel_radius
	config.temperature = args.temperature

	print("=" * 70)
	print("PALIMPSESTE — conversational chat training")
	print("=" * 70)
	print(f"config: D={config.D:,} context_window={config.context_window} "
	f"kernel_radius={config.kernel_radius} temperature={config.temperature}")
	print(f"Q/A pairs: {len(pairs)}")

	rng = np.random.default_rng(args.seed)
	lm = HFPalimpsesteLM(config=config, rng=rng)

	# build vocab from all questions + answers
	full_text = "".join(q + a for q, a in pairs)
	lm.build_tokenizer(full_text)
	print(f"vocab: {lm.tokenizer.vocab_size} chars")

	# train (O(1) per token, no gradient)
	import time
	total = 0
	for epoch in range(args.epochs):
	t0 = time.perf_counter()
	n = lm.train_on_qa_pairs(pairs, verbose=(args.epochs == 1))
	dt = time.perf_counter() - t0
	total += n
	print(f" epoch {epoch+1}/{args.epochs}: {n} tokens in {dt:.2f}s "
	f"({n/max(dt,1e-9):,.0f} tok/s), \|M\|={len(lm.mem)}")

	print(f"\ntotal: {total} tokens, \|M\| = {len(lm.mem):,}")

	# test responses on a few sample questions using the Conversation layer
	print("\n--- sample responses ---")
	from palimseste.chat import Conversation
	conv = Conversation(model=lm, learn_live=True)
	conv.register_questions(pairs)
	test_inputs = [
	"bonjour",
	"qui es-tu",
	"comment tu apprends",
	"tu utilises un gpu",
	"tu utilise un gpu", # typo
	"who are you",
	"how do you learn",
	"c quoi palimpseste",
	"merci",
	]
	for q in test_inputs:
	resp = conv.respond(q, temperature=0.0, seed=0)
	print(f" Q: {q}")
	print(f" A: {resp}")
	print()
	conv.reset() # don't save conversation history into the model

	# save
	out_dir = Path(args.output)
	print(f"saving to {out_dir} ...")
	lm.save_pretrained(out_dir)

	# push to hub
	if args.push_to_hub:
	print(f"\npushing to HF Hub: {args.push_to_hub} ...")
	url = lm.push_to_hub(args.push_to_hub, token=args.hub_token)
	print(f"pushed: {url}")

	print("\n" + "=" * 70)
	print("done. Chat with:")
	print(f" python examples/chat.py --model {out_dir}")
	print("=" * 70)


	if __name__ == "__main__":
	main()