wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-08-16 11:14:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

[feature] working trainer code

Browse Source
This commit is contained in:
theblackcat102
2022-12-31 03:02:10 +00:00
parent ad98a28241
commit bcd5c52b3b
6 changed files with 174 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/reward/instructor/tests/test_dataset.py
+5 -5
View File
@@ -1,26 +1,26 @@
from transformers import AutoTokenizer
from torch.utils.data import DataLoader
from rank_datasets import WebGPT, HFSummary, CollateFN
from rank_datasets import WebGPT, HFSummary, DataCollatorForMultipleChoice
def test_hfsummary():
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
collate_fn = CollateFN(tokenizer)
collate_fn = DataCollatorForMultipleChoice(tokenizer, max_length=200)
dataset = HFSummary()
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
for batch in dataloader:
print(batch[0]['input_ids'].shape)
print(batch['input_ids'].shape)
def test_webgpt():
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
collate_fn = CollateFN(tokenizer)
collate_fn = DataCollatorForMultipleChoice(tokenizer, max_length=200)
dataset = WebGPT()
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
for batch in dataloader:
print(batch[0]['input_ids'].shape)
print(batch['input_ids'].shape)
if __name__ == "__main__":