fix: handle tool messages in last round loss scaling

[Bakuma-sea]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Fix incorrect last_round loss scaling for agent conversations containing tool responses.

Previously, get_last_user_round only recognized messages with the user role. In tool-call conversations, a tool message starts the final query-response round, so the last round could be identified incorrectly and loss scaling could be applied to earlier assistant responses.

This change treats both user and tool messages as valid round-start messages. It also adds a regression test for a conversation ending with a tool response followed by the final assistant response.

Experiment results

• Added regression test: tests/loss_scale/test_last_user_round.py
• git diff --check passed.

[tastelikefeet]

Hi, can you lint the code via 'pre-commit run --all-files' please?

[Bakuma-sea]

Hi, can you lint the code via 'pre-commit run --all-files' please?

Thanks for the reminder. I ran pre-commit run --all-files, applied the yapf formatting changes, and all hooks pass now.