0
$\begingroup$

I would like to train test split a list of texts with the associated entities so there are no entities overlapping splits.

Ensuring no overlaps is challenging: I currently achieve it with 2 groupby operations through apache beam. I was wondering if there was a faster or more scalable way of doing this train test split?

Another way of formulating the question is I am looking for an efficient way to calculate connected components of a graph.

I have about 7 million documents with 250,000 entities and this dataset is likely to grow in the future.

INPUT

ENTITIES TEXT
e1       TextA
e1, e2   TextB
e3       TextC

I would like to have the outputs:

TRAIN SPLIT

ENTITIES TEXT
e1       TextA
e1, e2   TextB

TEST SPLIT

ENTITIES TEXT
e3       TextC

MY APPROACH

Initial groupby entities:

e1 [{"text":"TextA", "entities":["e1"]}, {"text":"TextB", "entities":["e1","e2"]}]
e2 [{"text":"TextB", "entities":["e1","e2"]}]
e3 [{"text":"TextC", "entities":["e3"]}]

Next I create a cooccuring entity key:

e1-e2 {"text":"TextA", "entities":["e1"]}
e1-e2 {"text":"TextB", "entities":["e1","e2"]}
e1-e2 {"text":"TextB", "entities":["e1","e2"]}
e3 {"text":"TextC", "entities":["e3"]}

I then groupby on this cooccuring key:

e1-e2 [{"text":"TextA", "entities":["e1"]}, {"text":"TextB", "entities":["e1","e2"]}]
e3 [{"text":"TextC", "entities":["e3"]}]

Finally I train test split these grouped entities.

Sadly my approach doesn't scale well so I am looking for alternatives.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.