Pass several input fields, only one field would get optimized

[thangld201]

Hi @zou-group, thank you very much for your works!

I am working on a problem where (i) the input contains multiple fields, only one field general instruction would remain constant and defined by us (to be optimized), whereas the other fields are specified by the user at inference time (dataset-dependent). Also, (ii) this task requires structured output so I'm also preparing few-shot examples for it (must-have to teach models to output the right format).

I haven't seen examples on how to (i) pass several input fields (one requires gradient only) to the model and loss function or (ii) pass few shot examples. I have gone through simple use cases such as question answering, but no sure how to handle these two problems (i) and (ii). Could you help me (e.g. where to start, besides the notebook examples) ?

My current codes look like this (which run but results are terrible):

import textgrad as tg
# each example has 4 fields: ["general_instruction", "schema", "input", "groundtruth"]
examples = read_from_file(example_path)
train_set = read_from_file(train_path)
engine = get_engine() # llm engine
model = tg.BlackboxLLM(engine)

evaluation_instruction = "Evaluate the generated JSON output. Provide concise and critical feedback."
evaluation_instruction = tg.Variable(evaluation_instruction,requires_grad=False,role_description="system prompt")
loss_fn = tg.TextLoss(evaluation_instruction)

constant_instruction = train_set[0]["general_instruction"]
constant_instruction = tg.Variable(constant_instruction, role_description="task instruction", requires_grad=True)
optimizer = tg.TGD(engine=engine,parameters=[constant_instruction])

epochs = 5
for _ in range(epochs):
    random_train_set = train_set.copy()
    rnd.shuffle(random_train_set)
    for sample in random_train_set:
        optimizer.zero_grad()

        cop_sample = deepcopy(sample)
        groundtruth = cop_sample["groundtruth"]
        cop_sample["instruction"] = constant_instruction.value
        cop_sample.pop("output")

        json_input = json.dumps(cop_sample,ensure_ascii=False)
        json_input = tg.Variable(json_input,requires_grad=False,role_description="JSON input")
        response = model(json_input)

        loss = loss_fn(response)
        loss.backward() 
        optimizer.step()

        # logging
        print(f"Current JSON input: {json_input.value}")
        print(f"Current instruction: {constant_instruction.value}")
        print(f"Expected output: {groundtruth}")
        print(f"Predicted output: {response}")
        print("-"*6)

cc @vinid @mertyg @nihalnayak @sugatoray @lupantech @ruanwz