wehub-resource-sync
336 lines
10 KiB
Python
336 lines
10 KiB
Python
"""
|
|
Trainer module tests
|
|
"""
|
|
|
|
import os
|
|
import unittest
|
|
import tempfile
|
|
|
|
import torch
|
|
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
|
|
|
from txtai.data import Data
|
|
from txtai.pipeline import HFTrainer, Labels, Questions, Sequences
|
|
|
|
|
|
class TestTrainer(unittest.TestCase):
|
|
"""
|
|
Trainer tests.
|
|
"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
"""
|
|
Create default datasets.
|
|
"""
|
|
|
|
cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
|
|
|
|
def testBasic(self):
|
|
"""
|
|
Test training a model with basic parameters
|
|
"""
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
|
|
|
|
labels = Labels((model, tokenizer), dynamic=False)
|
|
self.assertEqual(labels("cat")[0][0], 1)
|
|
|
|
def testCLM(self):
|
|
"""
|
|
Test training a model with causal language modeling
|
|
"""
|
|
|
|
trainer = HFTrainer()
|
|
|
|
# Test default parameters
|
|
model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation")
|
|
self.assertIsNotNone(model)
|
|
|
|
# Test pack merging
|
|
model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge="pack")
|
|
self.assertIsNotNone(model)
|
|
|
|
# Test no merging
|
|
model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge=None)
|
|
self.assertIsNotNone(model)
|
|
|
|
def testCustom(self):
|
|
"""
|
|
Test training a model with custom parameters
|
|
"""
|
|
|
|
# pylint: disable=E1120
|
|
model = AutoModelForSequenceClassification.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
|
|
tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer(
|
|
(model, tokenizer),
|
|
self.data,
|
|
self.data,
|
|
columns=("text", "label"),
|
|
do_eval=True,
|
|
output_dir=os.path.join(tempfile.gettempdir(), "trainer"),
|
|
)
|
|
|
|
labels = Labels((model, tokenizer), dynamic=False)
|
|
self.assertEqual(labels("cat")[0][0], 1)
|
|
|
|
def testDataFrame(self):
|
|
"""
|
|
Test training a model with a mock pandas DataFrame
|
|
"""
|
|
|
|
class TestDataFrame:
|
|
"""
|
|
Test DataFrame
|
|
"""
|
|
|
|
def __init__(self, data):
|
|
# Get list of columns
|
|
self.columns = list(data[0].keys())
|
|
|
|
# Build columnar data view
|
|
self.data = {}
|
|
for column in self.columns:
|
|
self.data[column] = Values([row[column] for row in data])
|
|
|
|
def __getitem__(self, column):
|
|
return self.data[column]
|
|
|
|
class Values:
|
|
"""
|
|
Test values list
|
|
"""
|
|
|
|
def __init__(self, values):
|
|
self.values = list(values)
|
|
|
|
def __getitem__(self, index):
|
|
return self.values[index]
|
|
|
|
def unique(self):
|
|
"""
|
|
Returns a list of unique values.
|
|
|
|
Returns:
|
|
unique list of values
|
|
"""
|
|
|
|
return set(self.values)
|
|
|
|
# Mock DataFrame
|
|
df = TestDataFrame(self.data)
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", df)
|
|
|
|
labels = Labels((model, tokenizer), dynamic=False)
|
|
self.assertEqual(labels("cat")[0][0], 1)
|
|
|
|
def testDataset(self):
|
|
"""
|
|
Test training a model with a mock Hugging Face Dataset
|
|
"""
|
|
|
|
class TestDataset(torch.utils.data.Dataset):
|
|
"""
|
|
Test Dataset
|
|
"""
|
|
|
|
def __init__(self, data):
|
|
self.data = data
|
|
self.unique = lambda _: [0, 1]
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
def __getitem__(self, index):
|
|
return self.data[index]
|
|
|
|
def column_names(self):
|
|
"""
|
|
Returns column names for this dataset
|
|
|
|
Returns:
|
|
list of columns
|
|
"""
|
|
|
|
return ["text", "label"]
|
|
|
|
# pylint: disable=W0613
|
|
def map(self, fn, batched, num_proc, remove_columns):
|
|
"""
|
|
Map each dataset row using fn.
|
|
|
|
Args:
|
|
fn: function
|
|
batched: batch records
|
|
|
|
Returns:
|
|
updated Dataset
|
|
"""
|
|
|
|
self.data = [fn(x) for x in self.data]
|
|
return self
|
|
|
|
ds = TestDataset(self.data)
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", ds)
|
|
|
|
labels = Labels((model, tokenizer), dynamic=False)
|
|
self.assertEqual(labels("cat")[0][0], 1)
|
|
|
|
def testEmpty(self):
|
|
"""
|
|
Test an empty training data object
|
|
"""
|
|
|
|
self.assertIsNone(Data(None, None, None).process(None))
|
|
|
|
def testKD(self):
|
|
"""
|
|
Test knowledge distillation
|
|
"""
|
|
|
|
# Base model
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
|
|
|
|
# Train with knowledge distillation
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data, teacher=(model, tokenizer))
|
|
|
|
labels = Labels((model, tokenizer), dynamic=False)
|
|
self.assertEqual(labels("cat")[0][0], 1)
|
|
|
|
def testMLM(self):
|
|
"""
|
|
Test training a model with masked language modeling.
|
|
"""
|
|
|
|
trainer = HFTrainer()
|
|
model, _ = trainer("hf-internal-testing/tiny-random-bert", self.data, task="language-modeling")
|
|
|
|
# Test model completed successfully
|
|
self.assertIsNotNone(model)
|
|
|
|
def testMultiLabel(self):
|
|
"""
|
|
Test training model with labels provided as a list
|
|
"""
|
|
|
|
data = []
|
|
for x in self.data:
|
|
data.append({"text": x["text"], "label": [0.0, 1.0] if x["label"] else [1.0, 0.0]})
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
|
|
|
|
labels = Labels((model, tokenizer), dynamic=False)
|
|
self.assertEqual(labels("cat")[0][0], 1)
|
|
|
|
def testPEFT(self):
|
|
"""
|
|
Test training a model with causal language modeling and PEFT
|
|
"""
|
|
|
|
trainer = HFTrainer()
|
|
model, _ = trainer(
|
|
"hf-internal-testing/tiny-random-gpt2",
|
|
self.data,
|
|
maxlength=16,
|
|
task="language-generation",
|
|
quantize=True,
|
|
lora=True,
|
|
)
|
|
|
|
# Test model completed successfully
|
|
self.assertIsNotNone(model)
|
|
|
|
def testQA(self):
|
|
"""
|
|
Test training a QA model
|
|
"""
|
|
|
|
# Training data
|
|
data = [
|
|
{"question": "What ingredient?", "context": "1 can whole tomatoes", "answers": "tomatoes"},
|
|
{"question": "What ingredient?", "context": "Crush 1 tomato", "answers": "tomato"},
|
|
{"question": "What ingredient?", "context": "1 yellow onion", "answers": "onion"},
|
|
{"question": "What ingredient?", "context": "Unwrap 2 red onions", "answers": "onions"},
|
|
{"question": "What ingredient?", "context": "1 red pepper", "answers": "pepper"},
|
|
{"question": "What ingredient?", "context": "Clean 3 red peppers", "answers": "peppers"},
|
|
{"question": "What ingredient?", "context": "1 clove garlic", "answers": "garlic"},
|
|
{"question": "What ingredient?", "context": "Unwrap 3 cloves of garlic", "answers": "garlic"},
|
|
{"question": "What ingredient?", "context": "3 pieces of ginger", "answers": "ginger"},
|
|
{"question": "What ingredient?", "context": "Peel 1 orange", "answers": "orange"},
|
|
{"question": "What ingredient?", "context": "1/2 lb beef", "answers": "beef"},
|
|
{"question": "What ingredient?", "context": "Roast 3 lbs of beef", "answers": "beef"},
|
|
{"question": "What ingredient?", "context": "1 pack of chicken", "answers": "chicken"},
|
|
{"question": "What ingredient?", "context": "Forest through the trees", "answers": None},
|
|
]
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data, data, task="question-answering", num_train_epochs=40)
|
|
|
|
questions = Questions((model, tokenizer), gpu=True)
|
|
self.assertTrue("onion" in questions(["What ingredient?"], ["Peel 1 onion"])[0])
|
|
|
|
def testRegression(self):
|
|
"""
|
|
Test training a model with a regression (continuous) output
|
|
"""
|
|
|
|
data = []
|
|
for x in self.data:
|
|
data.append({"text": x["text"], "label": x["label"] + 0.1})
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
|
|
|
|
labels = Labels((model, tokenizer), dynamic=False)
|
|
|
|
# Regression tasks return a single entry with the regression output
|
|
self.assertGreater(labels("cat")[0][1], 0.5)
|
|
|
|
def testRTD(self):
|
|
"""
|
|
Test training a language model with replaced token detection
|
|
"""
|
|
|
|
# Save directory
|
|
output = os.path.join(tempfile.gettempdir(), "trainer.rtd")
|
|
|
|
trainer = HFTrainer()
|
|
model, _ = trainer("hf-internal-testing/tiny-random-electra", self.data, task="token-detection", output_dir=output)
|
|
|
|
# Test model completed successfully
|
|
self.assertIsNotNone(model)
|
|
|
|
# Test output directories exist
|
|
self.assertTrue(os.path.exists(os.path.join(output, "generator")))
|
|
self.assertTrue(os.path.exists(os.path.join(output, "discriminator")))
|
|
|
|
def testSeqSeq(self):
|
|
"""
|
|
Test training a sequence-sequence model
|
|
"""
|
|
|
|
data = [
|
|
{"source": "Running again", "target": "Sleeping again"},
|
|
{"source": "Run", "target": "Sleep"},
|
|
{"source": "running", "target": "sleeping"},
|
|
]
|
|
|
|
trainer = HFTrainer()
|
|
model, tokenizer = trainer("t5-small", data, task="sequence-sequence", prefix="translate Run to Sleep: ", learning_rate=1e-3)
|
|
|
|
# Run run-sleep translation
|
|
sequences = Sequences((model, tokenizer))
|
|
result = sequences("translate Run to Sleep: run")
|
|
self.assertEqual(result.lower(), "sleep")
|