# A function that returns the log prob of each selected token
log_probs = calculate_log_probs(lm, generated_response)

buffer.extend([{
    "full_response": generated_response[i],
    "response_mask": response_mask[i], # A binary mask to denote which tokens in generated response are AI generated, 0 for system prompt and questions
    "old_log_probs": log_probs[i],
    "advantages": advantages[i]
} for i in range(len(generated_response))])