Machine learning for quantum states in embedding data