|
| 1 | +"""Unit tests for conversation fine-tuning and multi-turn chat handling.""" |
| 2 | + |
| 3 | +import tempfile |
| 4 | +import unittest |
| 5 | +from pathlib import Path |
| 6 | + |
| 7 | +import numpy as np |
| 8 | +import torch |
| 9 | +from tokenizers import Tokenizer |
| 10 | +from tokenizers.models import BPE |
| 11 | +from tokenizers.trainers import BpeTrainer |
| 12 | +from tokenizers.pre_tokenizers import ByteLevel |
| 13 | + |
| 14 | +from engine.config import ModelConfig |
| 15 | +from engine.data_core import _extract_structured_text |
| 16 | +from engine.microgpt_chat import STOP_SEQUENCES, MicroGPTChatSession |
| 17 | +from engine.target_masking import ( |
| 18 | + IGNORE_INDEX, |
| 19 | + InstructionDataset, |
| 20 | + find_completion_spans, |
| 21 | + mask_prompt_targets, |
| 22 | +) |
| 23 | + |
| 24 | + |
| 25 | +class TestConversationFineTune(unittest.TestCase): |
| 26 | + def setUp(self) -> None: |
| 27 | + self.tokenizer = Tokenizer(BPE(unk_token="<unk>")) |
| 28 | + self.tokenizer.pre_tokenizer = ByteLevel() |
| 29 | + trainer = BpeTrainer(special_tokens=["<unk>", "<|endoftext|>", "<pad>", "<bos>", "<eos>"], vocab_size=500) |
| 30 | + sample_corpus = [ |
| 31 | + "System: You are a conversational assistant.\n" |
| 32 | + "User: Hi\nAssistant: Hello! How can I help?\n" |
| 33 | + "User: Tell me more.\nAssistant: Certainly, here is more detail.\n<eos>\n", |
| 34 | + ] |
| 35 | + self.tokenizer.train_from_iterator(sample_corpus, trainer) |
| 36 | + |
| 37 | + def test_multi_turn_spans_detection(self) -> None: |
| 38 | + record = { |
| 39 | + "messages": [ |
| 40 | + {"role": "system", "content": "Be concise."}, |
| 41 | + {"role": "user", "content": "Question 1"}, |
| 42 | + {"role": "assistant", "content": "Answer 1"}, |
| 43 | + {"role": "user", "content": "Question 2"}, |
| 44 | + {"role": "assistant", "content": "Answer 2"}, |
| 45 | + ] |
| 46 | + } |
| 47 | + text = _extract_structured_text(record, "conversation") |
| 48 | + self.assertIn("System: Be concise.", text) |
| 49 | + self.assertIn("User: Question 1", text) |
| 50 | + self.assertIn("Assistant: Answer 1", text) |
| 51 | + |
| 52 | + spans = find_completion_spans(text) |
| 53 | + self.assertEqual(len(spans), 2) |
| 54 | + span1 = text[spans[0][0]:spans[0][1]] |
| 55 | + span2 = text[spans[1][0]:spans[1][1]] |
| 56 | + self.assertEqual(span1.strip(), "Answer 1") |
| 57 | + self.assertEqual(span2.strip(), "Answer 2") |
| 58 | + |
| 59 | + def test_multi_turn_target_masking(self) -> None: |
| 60 | + text = ( |
| 61 | + "System: Sys\n" |
| 62 | + "User: Q1\n" |
| 63 | + "Assistant: A1\n" |
| 64 | + "User: Q2\n" |
| 65 | + "Assistant: A2\n" |
| 66 | + "<eos>" |
| 67 | + ) |
| 68 | + enc = self.tokenizer.encode(text) |
| 69 | + targets = mask_prompt_targets(text, self.tokenizer, enc.ids, enc.offsets) |
| 70 | + self.assertEqual(len(targets), len(enc.ids)) |
| 71 | + |
| 72 | + # First tokens (System: Sys) must be -100 |
| 73 | + self.assertEqual(targets[0], IGNORE_INDEX) |
| 74 | + |
| 75 | + # Unmasked tokens must exist for both A1 and A2 |
| 76 | + unmasked = [tid for tid in targets if tid != IGNORE_INDEX] |
| 77 | + self.assertGreater(len(unmasked), 0) |
| 78 | + |
| 79 | + # Trailing <eos> must be unmasked |
| 80 | + self.assertEqual(targets[-1], enc.ids[-1]) |
| 81 | + |
| 82 | + def test_instruction_dataset_keeps_multi_turn_intact(self) -> None: |
| 83 | + # Single 5-turn conversation ending in EOS=4 |
| 84 | + tokens = np.array([10, 11, 12, 13, 14, 15, 16, 17, 4], dtype=np.int64) |
| 85 | + targets = np.array([-100, -100, 12, -100, -100, 15, 16, 17, 4], dtype=np.int64) |
| 86 | + ds = InstructionDataset(tokens, targets=targets, context_length=32, eos_token_id=4) |
| 87 | + # Should be exactly 1 discrete sample containing all turns |
| 88 | + self.assertEqual(len(ds), 1) |
| 89 | + x, y = ds[0] |
| 90 | + self.assertEqual(len(x), len(tokens) - 1) |
| 91 | + self.assertEqual(len(y), len(targets) - 1) |
| 92 | + |
| 93 | + def test_microgpt_chat_stop_sequences_contain_user(self) -> None: |
| 94 | + self.assertIn("\nUser:", STOP_SEQUENCES) |
| 95 | + self.assertIn("\nSystem:", STOP_SEQUENCES) |
| 96 | + self.assertIn("<eos>", STOP_SEQUENCES) |
| 97 | + |
| 98 | + def test_microgpt_chat_render_prompt_pruning(self) -> None: |
| 99 | + session = MicroGPTChatSession.__new__(MicroGPTChatSession) |
| 100 | + session.config = ModelConfig(vocab_size=500, context_length=64) |
| 101 | + session.tokenizer = self.tokenizer |
| 102 | + session._messages = [ |
| 103 | + {"role": "user", "content": "Very old question that should be dropped"}, |
| 104 | + {"role": "assistant", "content": "Very old answer that should be dropped"}, |
| 105 | + {"role": "user", "content": "Recent question"}, |
| 106 | + {"role": "assistant", "content": "Recent answer"}, |
| 107 | + ] |
| 108 | + prompt = "Current question" |
| 109 | + system_prompt = "Permanent System Instruction" |
| 110 | + |
| 111 | + rendered = session._render_prompt( |
| 112 | + prompt, system_prompt, reasoning_effort="None", thinking_enabled=False, max_tokens=16 |
| 113 | + ) |
| 114 | + |
| 115 | + # Permanent system instruction and latest turn MUST be preserved |
| 116 | + self.assertIn("System: Permanent System Instruction", rendered) |
| 117 | + self.assertIn("User: Current question", rendered) |
| 118 | + self.assertIn("Assistant:", rendered) |
| 119 | + |
| 120 | + # Token length must fit inside context_length - 16 |
| 121 | + enc = self.tokenizer.encode(rendered) |
| 122 | + self.assertLessEqual(len(enc.ids), session.config.context_length - 16) |
| 123 | + |
| 124 | + |
| 125 | +if __name__ == "__main__": |
| 126 | + unittest.main() |
0 commit comments