[P1] Convert reft model to hf model

[thu-yn]

I want to use some code data sets to fine-tune the basic model in the code task field. However, when I am evaluating, I hope to convert the generated reft model into a model format recognized by hf, so that it can directly call hf's generate function. Do you have any suggestions?

[frankaging]

hey @thu-yn thanks for your question!

if you train your model with our library, your resulting model should support generate function as:

instruction = "Which dog breed do people think is cuter, poodle or doodle?"

# tokenize and prepare the input
prompt = prompt_no_input_template % instruction
prompt = tokenizer(prompt, return_tensors="pt").to(device)

base_unit_location = prompt["input_ids"].shape[-1] - 1  # last position
_, reft_response = reft_model.generate(
    prompt, unit_locations={"sources->base": (None, [[[base_unit_location]]])},
    intervene_on_prompt=True, max_new_tokens=512, do_sample=True, 
    eos_token_id=tokenizer.eos_token_id, early_stopping=True
)
print(tokenizer.decode(reft_response[0], skip_special_tokens=True))

all of our evaluation scripts rely on huggingface's generate API --- so it should just work out of the box! let me know if you find any issue with that! thanks! note that the generate takes a couple more arguments, e.g., the intervening token positions. internally, we call huggingface generate call (which is done in our pyvene library https://github.com/stanfordnlp/pyvene/blob/main/pyvene/models/intervenable_base.py#L1585)