Train a small context LM

import whileai as wai

env = wai.methods.KVLog()
task = env.tasks(1)[0]
print(env.messages(task, 0, "")[1]["content"])  # the empty file and the first chunk
print(task["question"], task["gold"])

import whileai as wai

clm = wai.methods.ContextFile()  # w_eff 0.25, the paper's run
# Four trajectories: three right at costs 900, 1100 and 600, one wrong.
for steps in clm.credit([1, 1, 1, 0], [900, 1100, 600, 800], edits=5):
    print([round(a, 2) for a in steps])  # five edits, then the answer

import whileai as wai
from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer

env = wai.methods.KVLog()
clm = wai.methods.ContextFile()
rows = [{**t, "prompt": env.messages(t, 0, "")} for t in env.tasks(512)]
cfg = GRPOConfig(num_generations=8, num_iterations=1, beta=0.0, max_completion_length=256)
trainer = clm.trainer(GRPOTrainer)(
    model="Qwen/Qwen2.5-1.5B-Instruct",
    reward_funcs=[lambda completions, **_: [0.0] * len(completions)],  # the trainer pays the credit
    args=cfg,
    train_dataset=Dataset.from_list(rows),
    env=env,
)
trainer.train()