Skip to content

Fix tests #7

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

Draft
wants to merge 6 commits into
base: main
Choose a base branch
from
Draft

Fix tests #7

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
29f53aa
update
bcui19 Feb 20, 2025
14c4741
merge
bcui19 Feb 20, 2025
198c7ae
fix
bcui19 Feb 20, 2025
5074ae2
update
bcui19 Feb 20, 2025
5ae3324
merge
bcui19 May 19, 2025
f47cdc4
update
bcui19 May 19, 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
36 changes: 26 additions & 10 deletions tests/test_reward_learning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ def gen_random_batch(
test_cfg: Union[DictConfig, ListConfig],
inputs: Optional[list[str]] = None,
):
dataloader_type = test_cfg['train_loader']['name'] # type: ignore
# inputs can be [], ['input_ids'], ['input_ids', 'inputs_embeds'], and ['inputs_embeds']
# default to only input ids
if inputs == None:
Expand All @@ -145,18 +146,30 @@ def gen_random_batch(
size=(batch_size, test_cfg.max_seq_len * 2),
dtype=torch.int64,
).to(device)
batch['chosen_len'] = (
torch.ones(
size=(batch_size,),
if dataloader_type == 'pairwise_preference':
batch['chosen_len'] = (
torch.ones(
size=(batch_size,),
dtype=torch.int64,
) * test_cfg.max_seq_len
).to(device)
batch['rejected_len'] = (
torch.ones(
size=(batch_size,),
dtype=torch.int64,
) * test_cfg.max_seq_len
).to(device)
elif dataloader_type == 'finegrained_preference':
batch['text'] = torch.ones(
size=(batch_size, test_cfg.max_seq_len * 2),
dtype=torch.int64,
) * test_cfg.max_seq_len
).to(device)
batch['rejected_len'] = (
torch.ones(
).to(device)
batch['labels'] = torch.zeros((batch_size, 1),
dtype=torch.float32).to(device)
batch['text_len'] = torch.ones(
size=(batch_size,),
dtype=torch.int64,
) * test_cfg.max_seq_len
).to(device)
).to(device) * test_cfg.max_seq_len
return batch


Expand Down Expand Up @@ -248,7 +261,10 @@ def test_hf_train(
):
model_type, dataset_cls, collate_fn = model_params
model_name = 'jdchang/llama3-small'
tokenizer = AutoTokenizer.from_pretrained(model_name, pad_token='[PAD]')
tokenizer = AutoTokenizer.from_pretrained(
model_name,
pad_token='<|reserved_special_token_0|>',
)
max_seq_len = 10

dataset = dataset_cls(size=32, max_seq_len=max_seq_len)
Expand Down
Loading