Skip to content

Add Support For Tools and Documents Chat Template #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fabianlim merged 3 commits into from
Jun 18, 2025
Merged

Add Support For Tools and Documents Chat Template #3

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d639780
pass tools and docs
fabianlim Jun 11, 2025
30e3039
code improvement
fabianlim Jun 11, 2025
ecd090c
minor fix
fabianlim Jun 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion open_instruct/finetune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
from accelerate.utils import InitProcessGroupKwargs, set_seed
from datasets import load_dataset
from huggingface_hub import HfApi
from padding_free_collator import TensorDataCollatorWithFlattening
from .padding_free_collator import TensorDataCollatorWithFlattening
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
Expand Down Expand Up @@ -508,6 +508,12 @@ def encode_sft_example(example, tokenizer, max_seq_length):
We use the `apply_chat_template` function from the tokenizer to tokenize the messages and prepare the input and label tensors.
"""
messages = example["messages"]

additional_inputs = {}
for k in ["tools", "documents"]:
if k in example:
additional_inputs[k] = example[k]

if len(messages) == 0:
raise ValueError("messages field is empty.")
input_ids = tokenizer.apply_chat_template(
Expand All @@ -518,6 +524,7 @@ def encode_sft_example(example, tokenizer, max_seq_length):
truncation=True,
max_length=max_seq_length,
add_generation_prompt=False,
**additional_inputs,
)
labels = input_ids.clone()
# mask the non-assistant part for avoiding loss
Expand All @@ -537,6 +544,7 @@ def encode_sft_example(example, tokenizer, max_seq_length):
truncation=True,
max_length=max_seq_length,
add_generation_prompt=False,
**additional_inputs,
).shape[1]
# next, we calculate the end index of this non-assistant message
if (
Expand All @@ -554,6 +562,7 @@ def encode_sft_example(example, tokenizer, max_seq_length):
truncation=True,
max_length=max_seq_length,
add_generation_prompt=True,
**additional_inputs,
).shape[1]
else:
# for the last message or the message that doesn't follow with an assistant message,
Expand All @@ -566,6 +575,7 @@ def encode_sft_example(example, tokenizer, max_seq_length):
truncation=True,
max_length=max_seq_length,
add_generation_prompt=False,
**additional_inputs,
).shape[1]
# set the label to -100 for the non-assistant part
labels[:, message_start_idx:message_end_idx] = -100
Expand Down