Remove unnecessary new test file and consolidate test into existing class

This commit is contained in:
davidmezzetti
2026-09-29 07:10:09 -04:00
parent 95d4c1c881
commit 45d334f4ce
2 changed files with 19 additions and 40 deletions
@@ -1,39 +0,0 @@
"""
Texts module tests
"""
import unittest
from transformers import AutoTokenizer
from txtai.data import Texts
class TestTexts(unittest.TestCase):
"""
Texts tests.
"""
@classmethod
def setUpClass(cls):
"""
Create tokenizer.
"""
cls.tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
def testPack(self):
"""
Test packing rows into chunks up to maxlength
"""
rows = ["a b c d", "e f g h", "i j k l"]
length = len(self.tokenizer(rows[0])["input_ids"])
# Two rows fill maxlength exactly and are packed into one chunk
packed = Texts(self.tokenizer, None, length * 2, "pack").process({"text": list(rows)})
self.assertEqual([len(chunk) for chunk in packed["input_ids"]], [length * 2, length])
# Rows are never split across chunks
packed = Texts(self.tokenizer, None, length * 2 - 1, "pack").process({"text": list(rows)})
self.assertEqual([len(chunk) for chunk in packed["input_ids"]], [length] * 3)
@@ -15,7 +15,7 @@ import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from txtai.data import Data, Questions as QuestionsData
from txtai.data import Data, Questions as QuestionsData, Texts
from txtai.models import Lemur, Models, PoolingFactory
from txtai.pipeline import HFTrainer, Labels, LemurTrainer, Questions, Sequences
@@ -691,6 +691,24 @@ class TestTrainer(unittest.TestCase):
labels = Labels((model, tokenizer), dynamic=False)
self.assertEqual(labels("cat")[0][0], 1)
def testPack(self):
"""
Test packing rows into chunks up to maxlength
"""
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
rows = ["a b c d", "e f g h", "i j k l"]
length = len(tokenizer(rows[0])["input_ids"])
# Two rows fill maxlength exactly and are packed into one chunk
packed = Texts(tokenizer, None, length * 2, "pack").process({"text": list(rows)})
self.assertEqual([len(chunk) for chunk in packed["input_ids"]], [length * 2, length])
# Rows are never split across chunks
packed = Texts(tokenizer, None, length * 2 - 1, "pack").process({"text": list(rows)})
self.assertEqual([len(chunk) for chunk in packed["input_ids"]], [length] * 3)
def testPEFT(self):
"""
Test training a model with causal language modeling and PEFT