Choose Few-Shot Examples by Similarity and Diversity

Few-shot examples work best when they're similar to the user's input but not near-duplicates of each other. Select k examples from a pool of embeddings.

Implement select_examples(query, pool, k, max_sim=0.95, label_balance=None):

• Rank pool items by cosine similarity to the query (highest first, ties by index).
• Walk down the ranking and skip any candidate whose similarity to an already selected example exceeds max_sim.
• If label_balance is given (a list of labels, one per pool item), don't select more than ceil(k / number_of_distinct_labels) examples from any one label.
• Return the selected indices, in selection order. Return fewer than k if the constraints run out.