add context training and testing