`attribute_annotate.py` is not worked by KeyError: 'exceeded'

[AtsunoriFujita]

Describe the bug

When running attribute_annotate.py, an error occurs in get_reward(). This is due to KeyError: 'exceeded'.

Traceback (most recent call last):
  File "/opt/NeMo-Aligner/examples/nlp/data/steerlm/attribute_annotate.py", line 150, in <module>
    main(prepare_args())
  File "/opt/NeMo-Aligner/examples/nlp/data/steerlm/attribute_annotate.py", line 96, in main
    _ = get_reward(["hello world!"], host=args.host, port=args.port, model_name=args.model_name)
  File "/opt/NeMo-Aligner/examples/nlp/data/steerlm/attribute_annotate.py", line 74, in get_reward
    all_exceeded += output_dict["exceeded"].tolist()
KeyError: 'exceeded'

Steps/Code to reproduce bug Referring to the following https://docs.nvidia.com/nemo-framework/user-guide/latest/modelalignment/steerlm.html#step-4-generate-annotations

export MODEL="/workspace/models/Llama3-70B-SteerLM-RM"

python /opt/NeMo-Aligner/examples/nlp/gpt/serve_reward_model.py \
    rm_model_file=${MODEL} \
    trainer.num_nodes=1 \
    trainer.devices=8 \
    ++model.tensor_model_parallel_size=4 \
    ++model.pipeline_model_parallel_size=1 \
    inference.inference_micro_batch_size=8 \
    inference.port=1424

python /opt/NeMo-Aligner/examples/nlp/data/steerlm/attribute_annotate.py \
      --input-file=data/oasst/train.jsonl \
      --output-file=data/oasst/train_labeled.jsonl \
      --port=1424

Expected behavior

There is no error if there is or is not an "exceeded" key in the inference output.

Environment overview (please complete the following information)

• DGX-C A100 * 8
• nvcr.io/nvidia/nemo:24.07

Environment details

If NVIDIA docker image is used you don't need to specify these. Otherwise, please provide:

• OS version
• PyTorch version
• Python version

Additional context

Add any other context about the problem here. Example: GPU model