Here are the steps to achieve goal:

• Make all modality specific logic only in Dataset and Model. Rest of the classes have to be modality agnostic!
• Support texts
• Stop using keys to accumulate modality related information. Such information has to be processed inside Dataset.
• Introduce RetrievalResults class to store retrieval results. It has to work with any modality. Also: optimize memory by avoiding storing full distances matrix. Implement visualisation logic here to have possibility to draw something outside of EmbeddingMetrics (the modality specific logic of drawing a particular item has to be implemented in Dataset).