update training test file

jshuadvd · Jun 10, 2024 · bc4a9c8 · bc4a9c8
1 parent 15d18d7
commit bc4a9c8
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/tests/test_train.py b/tests/test_train.py
@@ -19,3 +19,16 @@ def test_collate_fn():
     inputs, targets = collate_fn(batch)
     assert inputs.shape == (2, 3)
     assert targets.shape == (2, 3)
+
+
+def test_preprocess_data():
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    data = "This is a test."
+    sequences = preprocess_data(data, tokenizer, max_length=10, overlap=5)
+    assert len(sequences) > 0
+
+
+def test_validate_targets():
+    targets = [[1, 2, 3], [4, 5, 6]]
+    vocab_size = 10
+    assert validate_targets(targets, vocab_size) == True