Upload tests/test_chat.py with huggingface_hub

8f2c86e verified 4 months ago

5.93 kB

"""Tests for conversational Q/A training and respond() (``palimseste.lm``).
Verifies:
  - train_on_qa_pairs writes traces to M (O(1) per token)
  - respond() returns a non-empty string after training
  - respond() recovers the exact trained answer for a seen question (greedy)
  - respond() handles unseen questions gracefully (no crash)
  - the episode format (BOS q EOS BOS a EOS) is consistent
  - generate_from_ids respects stop_on_eos
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.tokenizer import BOS, EOS


def _lm(D=5000, radius=30, ctx_window=48):
    cfg = PalimpsesteConfig(D=D, context_window=ctx_window, kernel_radius=radius,
                            temperature=0.3)
    return PalimpsesteForCausalLM(config=cfg)


# ----------------------------------------------------------------- train
def test_train_on_qa_pairs_grows_memory():
    lm = _lm()
    pairs = [("hello", "hi there"), ("bye", "goodbye")]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    n = lm.train_on_qa_pairs(pairs)
    assert n > 0
    assert len(lm.mem) == n


def test_train_on_qa_pairs_requires_tokenizer():
    lm = _lm()
    with pytest.raises(AssertionError):
        lm.train_on_qa_pairs([("q", "a")])


# ----------------------------------------------------------------- respond
def test_respond_returns_exact_answer_for_seen_question():
    lm = _lm(radius=0)
    pairs = [
        ("bonjour", "salut je suis palimpseste"),
        ("qui es-tu", "je suis palimpseste"),
        ("comment tu apprends", "j'apprends en ecrivant"),
    ]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    # greedy: should recover the exact trained answer
    resp = lm.respond("bonjour", temperature=0.0, seed=0)
    assert resp == "salut je suis palimpseste"


def test_respond_recovers_multiple_answers():
    lm = _lm(radius=0)
    pairs = [
        ("hello", "hi i am palimpseste"),
        ("who are you", "i am palimpseste a hypervectorial cortex"),
        ("how do you learn", "i learn by writing to memory"),
        ("do you use a gpu", "no i do not use a gpu"),
    ]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    for q, expected in pairs:
        resp = lm.respond(q, temperature=0.0, seed=0)
        assert resp == expected, f"Q={q!r}: got {resp!r}, expected {expected!r}"


def test_respond_unseen_question_no_crash():
    lm = _lm()
    pairs = [("hello", "hi there")]
    full = "hellohi there"
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    # unseen question should not crash; may return empty or partial
    resp = lm.respond("what is the weather", temperature=0.5, seed=0)
    assert isinstance(resp, str)


def test_respond_empty_input_no_crash():
    lm = _lm()
    lm.build_tokenizer("abc")
    lm.train_on_qa_pairs([("a", "b")])
    resp = lm.respond("", temperature=0.0, seed=0)
    assert isinstance(resp, str)


def test_respond_respects_max_tokens():
    lm = _lm(radius=0)
    # long answer
    pairs = [("q", "a" * 100)]
    full = "q" + "a" * 100
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    resp = lm.respond("q", max_new_tokens=5, temperature=0.0, seed=0)
    assert len(resp) <= 5


# ----------------------------------------------------------------- generate_from_ids
def test_generate_from_ids_stop_on_eos():
    lm = _lm(radius=0)
    text = "hello world"
    lm.build_tokenizer(text)
    lm.train_on_text(text)
    ids = lm.tokenizer.encode("hello", add_bos=True)
    out = lm.generate_from_ids(ids, max_new_tokens=50, temperature=0.0, seed=0,
                               stop_on_eos=True)
    assert isinstance(out, str)


def test_generate_from_ids_cold_start():
    lm = _lm()
    lm.build_tokenizer("abc")
    # untrained context -> cold start -> empty string
    out = lm.generate_from_ids([BOS], max_new_tokens=10, temperature=0.0, seed=0)
    assert isinstance(out, str)


# ----------------------------------------------------------------- episode format
def test_qa_episode_format_consistent():
    """The episode BOS q EOS BOS a EOS must be used consistently in
    train_on_qa_pairs and respond()."""
    lm = _lm(radius=0)
    pairs = [("hi", "hello there")]
    full = "hihello there"
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    # respond() must encode the question the same way: BOS q EOS BOS
    # then generate. The first generated token should be 'h' (start of "hello there")
    q_ids = lm.tokenizer.encode("hi", add_bos=True, add_eos=True)
    ctx = q_ids + [BOS]
    logits = lm._logits(ctx)
    import numpy as np
    top = int(np.argmax(logits))
    # the first answer char is 'h'
    assert lm.tokenizer.id2char[top] == "h"


# ----------------------------------------------------------------- robustness
def test_qa_with_french_accents():
    lm = _lm(radius=0)
    pairs = [("comment ça va", "ça va bien merci")]
    full = "comment ça vaça va bien merci"
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    resp = lm.respond("comment ça va", temperature=0.0, seed=0)
    assert "ça va bien" in resp or resp == "ça va bien merci"


def test_qa_multiple_pairs_no_interference():
    """Training on multiple Q/A pairs should not cause one answer to bleed
    into another (each episode is independent with its own BOS/EOS boundary)."""
    lm = _lm(radius=0)
    pairs = [
        ("aaa", "xxx"),
        ("bbb", "yyy"),
        ("ccc", "zzz"),
    ]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    assert lm.respond("aaa", temperature=0.0, seed=0) == "xxx"
    assert lm.respond("bbb", temperature=0.0, seed=0) == "yyy"
    assert lm.respond("ccc", temperature=0.0, seed=0) == "zzz"

	"""Tests for conversational Q/A training and respond() (``palimseste.lm``).

	Verifies:
	- train_on_qa_pairs writes traces to M (O(1) per token)
	- respond() returns a non-empty string after training
	- respond() recovers the exact trained answer for a seen question (greedy)
	- respond() handles unseen questions gracefully (no crash)
	- the episode format (BOS q EOS BOS a EOS) is consistent
	- generate_from_ids respects stop_on_eos
	"""

	from __future__ import annotations

	import numpy as np
	import pytest

	from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
	from palimseste.tokenizer import BOS, EOS


	def _lm(D=5000, radius=30, ctx_window=48):
	cfg = PalimpsesteConfig(D=D, context_window=ctx_window, kernel_radius=radius,
	temperature=0.3)
	return PalimpsesteForCausalLM(config=cfg)


	# ----------------------------------------------------------------- train
	def test_train_on_qa_pairs_grows_memory():
	lm = _lm()
	pairs = [("hello", "hi there"), ("bye", "goodbye")]
	full = "".join(q + a for q, a in pairs)
	lm.build_tokenizer(full)
	n = lm.train_on_qa_pairs(pairs)
	assert n > 0
	assert len(lm.mem) == n


	def test_train_on_qa_pairs_requires_tokenizer():
	lm = _lm()
	with pytest.raises(AssertionError):
	lm.train_on_qa_pairs([("q", "a")])


	# ----------------------------------------------------------------- respond
	def test_respond_returns_exact_answer_for_seen_question():
	lm = _lm(radius=0)
	pairs = [
	("bonjour", "salut je suis palimpseste"),
	("qui es-tu", "je suis palimpseste"),
	("comment tu apprends", "j'apprends en ecrivant"),
	]
	full = "".join(q + a for q, a in pairs)
	lm.build_tokenizer(full)
	lm.train_on_qa_pairs(pairs)
	# greedy: should recover the exact trained answer
	resp = lm.respond("bonjour", temperature=0.0, seed=0)
	assert resp == "salut je suis palimpseste"


	def test_respond_recovers_multiple_answers():
	lm = _lm(radius=0)
	pairs = [
	("hello", "hi i am palimpseste"),
	("who are you", "i am palimpseste a hypervectorial cortex"),
	("how do you learn", "i learn by writing to memory"),
	("do you use a gpu", "no i do not use a gpu"),
	]
	full = "".join(q + a for q, a in pairs)
	lm.build_tokenizer(full)
	lm.train_on_qa_pairs(pairs)
	for q, expected in pairs:
	resp = lm.respond(q, temperature=0.0, seed=0)
	assert resp == expected, f"Q={q!r}: got {resp!r}, expected {expected!r}"


	def test_respond_unseen_question_no_crash():
	lm = _lm()
	pairs = [("hello", "hi there")]
	full = "hellohi there"
	lm.build_tokenizer(full)
	lm.train_on_qa_pairs(pairs)
	# unseen question should not crash; may return empty or partial
	resp = lm.respond("what is the weather", temperature=0.5, seed=0)
	assert isinstance(resp, str)


	def test_respond_empty_input_no_crash():
	lm = _lm()
	lm.build_tokenizer("abc")
	lm.train_on_qa_pairs([("a", "b")])
	resp = lm.respond("", temperature=0.0, seed=0)
	assert isinstance(resp, str)


	def test_respond_respects_max_tokens():
	lm = _lm(radius=0)
	# long answer
	pairs = [("q", "a" * 100)]
	full = "q" + "a" * 100
	lm.build_tokenizer(full)
	lm.train_on_qa_pairs(pairs)
	resp = lm.respond("q", max_new_tokens=5, temperature=0.0, seed=0)
	assert len(resp) <= 5


	# ----------------------------------------------------------------- generate_from_ids
	def test_generate_from_ids_stop_on_eos():
	lm = _lm(radius=0)
	text = "hello world"
	lm.build_tokenizer(text)
	lm.train_on_text(text)
	ids = lm.tokenizer.encode("hello", add_bos=True)
	out = lm.generate_from_ids(ids, max_new_tokens=50, temperature=0.0, seed=0,
	stop_on_eos=True)
	assert isinstance(out, str)


	def test_generate_from_ids_cold_start():
	lm = _lm()
	lm.build_tokenizer("abc")
	# untrained context -> cold start -> empty string
	out = lm.generate_from_ids([BOS], max_new_tokens=10, temperature=0.0, seed=0)
	assert isinstance(out, str)


	# ----------------------------------------------------------------- episode format
	def test_qa_episode_format_consistent():
	"""The episode BOS q EOS BOS a EOS must be used consistently in
	train_on_qa_pairs and respond()."""
	lm = _lm(radius=0)
	pairs = [("hi", "hello there")]
	full = "hihello there"
	lm.build_tokenizer(full)
	lm.train_on_qa_pairs(pairs)
	# respond() must encode the question the same way: BOS q EOS BOS
	# then generate. The first generated token should be 'h' (start of "hello there")
	q_ids = lm.tokenizer.encode("hi", add_bos=True, add_eos=True)
	ctx = q_ids + [BOS]
	logits = lm._logits(ctx)
	import numpy as np
	top = int(np.argmax(logits))
	# the first answer char is 'h'
	assert lm.tokenizer.id2char[top] == "h"


	# ----------------------------------------------------------------- robustness
	def test_qa_with_french_accents():
	lm = _lm(radius=0)
	pairs = [("comment ça va", "ça va bien merci")]
	full = "comment ça vaça va bien merci"
	lm.build_tokenizer(full)
	lm.train_on_qa_pairs(pairs)
	resp = lm.respond("comment ça va", temperature=0.0, seed=0)
	assert "ça va bien" in resp or resp == "ça va bien merci"


	def test_qa_multiple_pairs_no_interference():
	"""Training on multiple Q/A pairs should not cause one answer to bleed
	into another (each episode is independent with its own BOS/EOS boundary)."""
	lm = _lm(radius=0)
	pairs = [
	("aaa", "xxx"),
	("bbb", "yyy"),
	("ccc", "zzz"),
	]
	full = "".join(q + a for q, a in pairs)
	lm.build_tokenizer(full)
	lm.train_on_qa_pairs(pairs)
	assert lm.respond("aaa", temperature=0.0, seed=0) == "xxx"
	assert lm.respond("bbb", temperature=0.0, seed=0) == "yyy"
	assert lm.respond("ccc", temperature=0.0, seed=0) == "zzz"