Implement rejection sampling for supervised fine-tuning data generation. Given a list of prompts where each prompt has K candidate outputs and their corresponding reward model scores, select the single best (highest-scoring) candidate output per prompt.
Implement rejection_sampling_best_of_k(candidates, scores) that:
candidates: a list of length N (number of prompts), where each element is a list of K candidate output strings.scores: a list of length N, where each element is a list of K floats representing the reward model scores for the corresponding candidates.N containing the selected candidate (the one with the highest score) for each prompt.If two or more candidates are tied for the highest score, return the one that appears first (lowest index). Assume K >= 1 for every prompt and that the inner lengths of candidates[i] and scores[i] match.