Text embeddings and data splitting

Question

I have created some document embeddings which were then used further in text classification tasks. After revisiting my code I was unsure about the workflow I used to train the document embeddings.

At the moment I am creating the document embeddings based on the complete corpus available at the time of training. After the training is done, I evaluate the model by looking whether it creates useful similarities between the document embeddings. Those embeddings are used then in machine learning models and that's where the embeddings will be split into train, test and validation sets.

Now my question is: Where is the right time to split the data? Should I do it before creating the document embeddings to prevent data leakage? I have used the mentioned approach because I viewed the creation of the document embeddings as a preprocessing step, so the computer can work with textual data. However, after I have put some thought into it, I think it's the wrong approach. I wanted to hear from more experienced NLP practitioners how they approach this task. Sorry for this very basic question.

Thanks.

Answer 1

TL;DR

If you are training the document-embedding model, then split the data before you convert the text into embeddings.

If you are using a pre-trained document-embedding model, then it won't matter and it is pre-processing step that it doesn't matter when you execute it.

Pipeline when training your own document-embedding model

1. Split your text data into train/validate/test sets.
2. Use your train set to train the document-embedding model.
3. Use your trained document-embedding model to convert train and validation sets to train your other model (e.g. classification model).
4. Test your final model by using your trained document-embedding model to convert the test set and test the trained final (classification) model.