[ML] Refactors the signature of the assignment map

[valeriy42]

• Have you signed the contributor license agreement?
• Have you followed the contributor guidelines?
• If submitting code, have you built your formula locally prior to submission with gradle check?
• If submitting code, is your pull request against main? Unless there is a good reason otherwise, we prefer pull requests against main and will backport as needed.
• If submitting code, have you checked that your submission is for an OS and architecture that we support?
• If you are submitting this code for a class then read our policy for that.