Bug fixes on document training. Merged with ui upstream.