Add back LM.forward()'s callback

[chenmoneygithub]

Revert #8098. This extra callback is useful to interpret the actual LM prompt and response. The callback on BaseLM.__call__ only captures the processed output.

[TomeHirata]

Shall we remove callback from BaseLM.__call__ then? It is counterintuitive to call the LM callback two times for each LM call, so we should change the flow a bit instead of just recovering the callback to LM class. Also, it was not consistent with DummyLM implementation.

[chenmoneygithub]

@TomeHirata Yes it's weird on the UI, so in fact I went through the same thinking process, and realized that both are useful.

BaseLM.__call__ has additional formatting logic from LM.forward(), so the traced outputs are different. I talked to Yuki about if we should make callbacks aware of the fn name besides the class name, but we realized that users may not understand the difference between LM.__call__ and LM.forward() as well.

[TomeHirata]

Got it, thank you for the context. I think the main problem here is that the separation between call and forward is not clear for users and it gets counterintuitive to call the on_lm_start callback twice. Is it possible to 1) add a method with (or rename forward to) a more explicit name like make_lm_call and 2) allow callback handlers to define another callback method like on_make_lm_call_start? We can mitigate the problem only with the latter but the former is a plus.

[chenmoneygithub]

@TomeHirata Yea that could help, but honestly I would not bother doing it because it turned out that users can manage to understand the double trace on the LM call either autonomously or with a simple guide when I talked to them. I will leave this to your call, if you feel double trace is too ugly, feel free to add a new hook and change MLflow side accordingly.

[chenmoneygithub]

@TomeHirata @okhat Btw, we do need to merge this for now. Deeplearning.ai is specifically asking us to add back the actual LM's input/output trace.