Machine learning approach for predicting set members Say I have a large training dataset containing sets of 40 items each, and each item in the set is unique (so every training input is a set $S=\{i_1, i_2, ..., i_{40}\}$), and there are more than 40 unique items that can be part of a set.

I would like to be able to predict which items are a probably member of a set, given some incomplete set. So let's take the following example:

Training data:
$S_1 = \{1,2,3\}$,
$S_2 = \{3,4,5\}$,
$S_3 = \{6,7,8\}$
Say I then have an input $S_4 = \{3\}$, I would want the approach to give back that 1, 2, 4 and 5 are more probable set members than 7, 8. Ideally with some probability value.

I've considered the following:
Using the apriori algorithm to learn some association rules. I wasn't sure how to interpret the support or lift as a probability of set membership.

Training a Multilayer perceptron on the input (probably one-hot encoded) to learn weights corresponding to the various input items. However, if I were to simply give the 40-item sets as input and output then the network would just learn to copy the input, giving no information about possible other set members. I've thought about giving all variations of the 40 item set as input, with the 40 item set as output, but this would result in $2^{40}$ possibilities per input which would be massive.

Is there some machine learning approach or data structure that could help in this situation?
You could train an embedding model. Each element would be projected onto a location a vector space based on its co-occurrence with other elements. Then finding similar elements could be done with a nearest neighbor search.