testing config

example_trainer/run_gsm8k_teacher_distill_single_terminal.sh

@@ -229,6 +229,7 @@ start_process "gsm8k_teacher_env" "${LOG_DIR}/env.log" \
     --env.teacher_base_url "http://localhost:${TEACHER_PORT}/v1" \
     --env.teacher_model_name "$TEACHER_MODEL" \
     --env.teacher_top_k "$TEACHER_TOP_K" \
+    --env.ensure_scores_are_not_same false \
     --openai.api_key "dummy" \
     --openai.base_url "http://localhost:${STUDENT_PORT}/v1" \
     --openai.model_name "$STUDENT_MODEL" \

example_trainer/vllm_api_server.py

@@ -360,6 +360,13 @@ async def _generate(request_dict: dict, raw_request: Request) -> Response:
         ret["prompt_token_ids"] = final_output.prompt_token_ids
         ret["token_ids"] = [x.token_ids for x in final_output.outputs]
 
+    if sampling_params.prompt_logprobs is not None and final_output.prompt_logprobs is not None:
+        ret["prompt_logprobs"] = [
+            {int(tok_id): lp.logprob for tok_id, lp in pos.items()}
+            if pos is not None else None
+            for pos in final_output.prompt_logprobs
+        ]
+
     logger.info(
         "POST /generate completed request_id=%s outputs=%s finish_reasons=%s",
         request_id,