TestTrainSplitter silently drops documents instead of adding to train set

MedCAT
1

Description

When using make_mc_train_test (or test_size > 0 in train_supervised_raw), the TestTrainSplitter.split() method seems to silently drops the majority of documents rather than assigning them to the train set. This results in a training set far smaller than expected.

Expected behaviour

With 472 annotated documents and test_size=0.1, approximately 425 documents should be assigned to the train set and ~47 to the test set.

Actual behaviour

Only ~20 documents appear in the train set. The remaining ~450 documents are assigned to neither train nor test set and are silently discarded.

Root cause

In medcat-v2/medcat/utils/data_utils.py, the split() method uses continue when the test quota is full, skipping documents entirely rather than routing them to the train set:

for i_document in np.random.permutation(range(0, num_of_docs)):
    if self.test_anns / self.total_anns >= self.test_size:
        continue  # document is silently dropped
    document = project['documents'][i_document]
    self._split_doc_train_test(document, cui_filter,
                               train_project, test_project)

Proposed fix

Add the document to the train set before continue:

for i_document in np.random.permutation(range(0, num_of_docs)):
    if self.test_anns / self.total_anns >= self.test_size:
        train_project['documents'].append(project['documents'][i_document])  # route to train
        continue
    document = project['documents'][i_document]
    self._split_doc_train_test(document, cui_filter,
                               train_project, test_project)

This appears to work well with 91 docs being assigned to test and 381 to train with a test_size=0.2; there is a good split of annotations as well, with 589 in the test set and 2931 in the train set, again with test_size=0.2.

Full disclosure, this was again diagnosed using claude AI.

Happy to create pull request if you feel this is the right kind of fix, although I note in the documentation for make_mc_train_test the comment ‘This is a disaster’ and not sure what other issues there are with it.

2

Thanks for the detailed report.

The issue with the test-train split is that it’s designed to try and do a smarter split than just a trivial one. And that makes the code paths more complicated.

With that said, what you’ve described sounds like it is indeed an issue and the solution looks sound as well.

So if you could do a PR for the fix, that would be greatly appreciated.

PS:
If you do a PR, please also include a test for this behaviour. I.e perhaps a synthetic dataset where you can (without the change) observe the same behaviour.