import unittest
import pandas as pd
from torch.utils.data import DataLoader
from transformers import BertTokenizer
from datareader.abstract_data_reader import InputExample
from medicalbert.datareader.chunked_data_reader import ChunkedDataReader
class TestChunkedDataReader(unittest.TestCase):
def __init__(self, methodName='runTest'):
super().__init__(methodName)
self.tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
self.config = {"max_sequence_length": 10, "target":"target", "num_sections": 10}
self.cdr = ChunkedDataReader(self.config, self.tokenizer)
def test_chunker_gen(self):
# create a test string
test_input = "Hi My name is Andrew Patterson and I made this".split()
from medicalbert.datareader.chunked_data_reader import ChunkedDataReader
sections = [section for section in ChunkedDataReader.chunks(test_input, 3)]
self.assertTrue(len(sections) == 4, "Correct number of sections returned")
self.assertEqual(sections[0], ['Hi', 'My', 'name'], "First section is correct")
self.assertEqual(sections[3], ['this'], "Last section is correct")
def assertInputFeatureIsValid(self, inputFeature, sep_index):
# check the length
self.assertTrue(len(inputFeature.input_ids) == 10)
self.assertTrue(len(inputFeature.segment_ids) == 10)
self.assertTrue(len(inputFeature.input_mask) == 10)
# now check that the cls and sep tokens are in the correct place.
self.assertTrue(inputFeature.input_ids[0] == 101)
self.assertTrue(inputFeature.input_ids[sep_index] == 102)
# now check that the padded space is filled with zeroes
expected = [0] * 501
actual = inputFeature.input_ids[11:]
self.assertTrue(expected, actual)
@staticmethod
def make_test_data():
# make a dummy dataset
examples_text = ["Hi My name is Andrew Patterson and I made this and so I must test it.",
"Hi My name is Andrew",
"Hi My name is Andrew Patterson and I made this and so I must test it.",
"Hi My name is Andrew", ]
examples_label = [1, 1, 1, 1]
data = {'text': examples_text, 'target': examples_label}
return pd.DataFrame.from_dict(data)
def test_build_fresh_dataset(self):
test_data = TestChunkedDataReader.make_test_data()
tensor_dataset = self.cdr.build_fresh_dataset(test_data)
self.assertTrue(4, len(tensor_dataset[0])) # This checks the number of features
print(tensor_dataset[0][0].shape)
def test_convert_section_to_feature_short(self):
# create a test string that is shorter than the max sequence length
test_input = "Hi My name is Andrew"
tokens = self.tokenizer.tokenize(test_input)
# convert to a feature
inputFeature = self.cdr.convert_section_to_feature(tokens, "1")
self.assertInputFeatureIsValid(inputFeature, 6)
def test_convert_section_to_feature_long(self):
# create a test string that is longer than the max sequence length
test_input = "Hi My name is Andrew Patterson and I made this and so I must test it."
tokens = self.tokenizer.tokenize(test_input)
# convert to a feature
inputFeature = self.cdr.convert_section_to_feature(tokens, "1")
self.assertInputFeatureIsValid(inputFeature, 9)
def test_convert_example_to_feature(self):
# create a test string that is longer than the max sequence length
test_input = "Hi My name is Andrew Patterson and I made this and so I must test it."
e = InputExample(None, test_input, None, 1)
result = self.cdr.convert_example_to_feature(e, 1)
if __name__ == '__main__':
unittest.main()
