Fix failing test for word model
Refactor trains helper with test