def build_single_shot_inference_calls()

in yourbench/utils/inference/inference_builders.py [0:0]


def build_single_shot_inference_calls(dataset, system_msg, stage_cfg, sampling_cfg):
    calls = []
    index_map = []

    for idx, row in enumerate(dataset):
        document_chunks = row.get("chunks") or []
        selected_chunks = sample_single_hop_chunks(document_chunks, sampling_cfg)

        for ch_idx, chunk in enumerate(selected_chunks):
            chunk_id = chunk.get("chunk_id", f"{idx}_{ch_idx}")
            chunk_text = chunk.get("chunk_text", "")
            user_msg = {
                "role": "user",
                "content": QUESTION_GENERATION_USER_PROMPT.format(
                    title=row.get("document_filename", f"doc_{idx}"),
                    document_summary=row.get("document_summary", ""),
                    text_chunk=chunk_text,
                    additional_instructions=stage_cfg.get("additional_instructions", ""),
                ),
            }
            calls.append(InferenceCall(messages=[system_msg, user_msg], tags=["single_shot_qa"]))
            index_map.append((idx, row.get("document_id", f"doc_{idx}"), chunk_id))

    return calls, index_map