Add multiple sampling support to lf.query and introduce lf.query_and_reduce.

[copybara-service]

Add multiple sampling support to lf.query and introduce lf.query_and_reduce.

This CL enhances lf.query to support parallel queries to multiple LLMs with multiple samples by accepting a list of models (lm) and the corresponding num_samples argument.

Additionally, this change introduces a new API, lf.query_and_reduce, which enables a map-reduce style workflow using lf.query. The motivation for lf.query_and_reduce is to maintain a clean 1:1 mapping at the interface level while providing a flexible mechanism for users to regularize or aggregate outputs using a reduce function.

Key Features:

• lf.query Enhancements:

  • Supports multiple language models queried in parallel.
  • Allows sampling multiple outputs from each model using the num_samples argument.

• lf.query_and_reduce:

  • A drop-in replacement for lf.query with compatible arguments and default behavior.
  • Provides a reduce function that processes multiple LLM outputs, enabling operations like majority voting, aggregation, or other forms of output post-processing.

Example usage: sampling Gpt4 and Gemini2 each with 5 samples and obtain the final output through majority voting.

lf.query_and_reduce(
  'compute 256 * 345',
  int,
  reduce=lambda answers: lf.query('Find majority from {{answers}}', int, lm=lf.llms.Gpt4(), answers=answers),
  lm=[lf.llms.Gpt4(), lf.llms.Gemini2()],
  num_samples=5,
  ...
)