[fix] update pairwise dataloader.

[Chen9154]

In forward() of reward_model.py (Line 62), if "chosen" and "rejected" are exactly the same, "inference" would turn to True, which shouldn't happen during the training procedure. However in class PairwiseDataset, "chosen" and "rejected" could be the same after truncation (this would be easily happen when prompts/posts are longer than max_length and we set padding_side = 'right'). So we filter out those cases from training data.

[jon-tow]

Hi, @Chen9154! Thanks for catching this edge case; looks good to me 👍

Would you be able to also make this same change to the test file for sake of completeness, please?

trlx/examples/summarize_rlhf/reward_model/gptj_reward_test.py

Lines 40 to 65 in 9fd1f0a

	class PairwiseDataset(Dataset):
	def __init__(self, pairs, tokenizer, max_length):
	self.chosen_input_ids = []
	self.chosen_attn_masks = []
	self.rejected_input_ids = []
	self.rejected_attn_masks = []
	for pair in pairs:
	chosen, rejected = pair["chosen"], pair["rejected"]
	chosen_encodings_dict = tokenizer(
	"<|startoftext|>" + chosen + "<|endoftext|>",
	truncation=True,
	max_length=max_length,
	padding="max_length",
	return_tensors="pt",
	)
	rejected_encodings_dict = tokenizer(
	"<|startoftext|>" + rejected + "<|endoftext|>",
	truncation=True,
	max_length=max_length,
	padding="max_length",
	return_tensors="pt",
	)
	self.chosen_input_ids.append(chosen_encodings_dict["input_ids"])
	self.chosen_attn_masks.append(chosen_encodings_dict["attention_mask"])
	self.rejected_input_ids.append(rejected_encodings_dict["input_ids"])
	self.rejected_attn_masks.append(rejected_encodings_dict["attention_mask"])

[Chen9154]

@jon-tow Thanks for the review! I have also made the same change to the test file.

[jon-tow]

Awesome! Thank you, @Chen9154!