Automate example selection in LLM few-shot learning

[sonichi]

Planning new features in autogen. Suggestions are welcome.

[qingyun-wu]

@jtongxin

[feiran-jia]

I have devised a preliminary solution for integrating the demonstration selection into autogen.

Create() Function
There is no need to modify the API of oai.ChatCompletion.create(). Users can incorporate the few-shot examples by writing a template function that defines how to add the examples and context. To illustrate this, I have provided an example below.

def few_shot_format(few_shot_set, cot_flag=False)->str:
    # This is a fun indicating the format of few-shot examples
    import guidance

    structure_program = guidance(
    ''' We first provide the problem statement, then the solution as the examples.
    ----
    {{~! display the few-shot examples ~}}
    {{~#each examples}}
    goal: {{this.goal}}
    sol1: {{this.sol1}}
    sol2: {{this.sol2}}
    label: {{this.label}}
    ---
    {{~/each}}
    ''')

    out = structure_program(
        examples=few_shot_set,
    )

    return str(out)

Then we write the template function:

def content_few_shot(context, train_data = None, **demonstration_params):
    
    prompt = "----\n Instruction: label: 0 refers to sol1 and 1 refers to sol2. Please just response 1 or 0:'  \n---\n"
    # Demonstration Selection Method
    few_shot_list = random_select(train_data, **demonstration_params) 
    # the template to intergrate the few-shot examples
    few_shot_prompt = few_shot_format(few_shot_list)
    problem = "\n".join(
        [
            "goal: "+ context["goal"],
            "sol1: "+ context["sol1"],
            "sol2:" + context["sol2"],
            "label: ",
        ]
    )
    return  "We will first show some examples: " + few_shot_prompt + prompt + problem 

demonstration_params = {'k':2}
template_few_shot_fn = partial(content_few_shot, train_data = train_data, **demonstration_params)

After we got the template function, we can perform create()

messages_few_shot = [
    {
        "role": "user",
        "content": template_few_shot_fn,
    },
]

response = oai.ChatCompletion.create(
    context = train_data[1],
    messages= messages_few_shot,
    config_list=[
        {
            "model": "gpt-3.5-turbo",
            "api_key": "<your key here>",
            "api_type": "open_ai",
            "api_base": "https://api.openai.com/v1",
            "api_version": None,
        },
    ],
)

[feiran-jia]

Tune Function

To incorporate the demonstration selection (or few-shot examples), we need to let the tune func know two things

• whether we are doing demonstration selection
• wether we are using the customized selection function

My thought is to add an argument called demonstration_params (a dict).

    demonstration_params = {
        'selection_methods':   '<function name>' or selection_func or a list of them
        '<hp_name>': '<hp>',  # the hyper parameters of the selection function.
        ...
    }

If demonstration_params is None, we just perform zero-shot completion. Otherwise, we do the few-shot learning (in-context learning) in several ways:

• demonstration_params['selection_methods'] is a built-in demonstration selection method. We will modify the message within the tune, and the user just need to provide a few-shot template without the demonstration selection. For example

  def few_shot_prompt_fun(context, few_shot_list = None)->str:
  # question and answer
  import guidance
  
  prompt = "---- Instruction: label -  0 refers to sol1 and 1 refers to sol2. Please just return one word 1 or 0:'  \n---\n"
  
  # define the guidance program
  if few_shot_list is not None:
     structure_program = guidance(
     ''' We first provide the problem statement, then the solution as the examples.
     ----
     {{~! display the few-shot examples ~}}
     {{~#each examples}}
     goal: {{this.goal}}
     sol1: {{this.sol1}}
     sol2: {{this.sol2}}
     label: {{this.label}}
     ---
     {{~/each}}
     ''')
  
     out = structure_program(
         examples=few_shot_list,
     )
     few_shot_prompt = str(out)
  else:
     few_shot_prompt = ""
  
  problem = "\n".join(
     [
         "goal: "+ context["goal"],
         "sol1: "+ context["sol1"],
         "sol2:" + context["sol2"],
         "label: ",
     ]
  )
  
  return few_shot_prompt + prompt + problem

• demonstration_params['selection_methods'] is a function. In this case, the user will provide both the few-shot template without the demonstration selection, and the selection function. Here provide the simple selection method

  def random_select(data, k = 10):
      import random
      selected_samples = random.sample(list(data), k)
      return selected_samples

• demonstration_params['selection_methods'] is a unknown str, and the selection function is not provided. In this case, the user should write the template function that defines how to add the examples and context. This should be similar to my example of create.

[RageAgainstTheMachine101]

@sonichi could you, please, tell are there any solutions for few-shot learning in autogen so far?

[sonichi]

Nothing in particular, except the templating which could be related. @qingyun-wu, anything to add?

[qingyun-wu]

@sonichi could you, please, tell are there any solutions for few-shot learning in autogen so far?

Thanks @sonichi and @RageAgainstTheMachine101. I am indeed considering adding a relevant feature and trying to find a good use cases. @RageAgainstTheMachine101 can you please briefly describe your use cases that need few-shot learning?

[RageAgainstTheMachine101]

@RageAgainstTheMachine101 can you please briefly describe your use cases that need few-shot learning?

@qingyun-wu yes, I can.

I am working with the following scenario: a requester asks some ad hoc question about users behavior in a product -> AI agent tries to map ad hoc question to SQL -> another AI agent tries to run this query against the db and fix errors if they occur -> the final result is provided to requester.

The main problem here is documentation knowledge. I resolved this problem with a system prompt with a lot of docs info.
Also, there is a problem with a special SQL dialect and code style. The second problem is resolved via a few shot learning approach very well. But the whole scenario was implemented without autogen.

[sonichi]

@RageAgainstTheMachine101 have you checked the teachable agent or retrieve chat in autogen? It sounds relevant from what you said.

[RageAgainstTheMachine101]

@RageAgainstTheMachine101 have you checked the teachable agent or retrieve chat in autogen? It sounds relevant from what you said.

Yes, I tried retrieve chat, but it consumed all available tokens per minute for an organization. I believe I should try teachable agent and find out how to limit tokens consuming.

Thank you!