mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-26 13:07:22 +08:00
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from transformers import AutoTokenizer
|
|
from torch.utils.data import DataLoader
|
|
from rank_datasets import WebGPT, HFSummary, DataCollatorForPairRank
|
|
from experimental_dataset import HFSummaryQuality, DataCollatorForSummaryScore
|
|
|
|
def test_hfsummary():
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
|
|
collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
|
|
dataset = HFSummary('train')
|
|
print(len(dataset))
|
|
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
|
|
for batch in dataloader:
|
|
batch['input_ids'].shape
|
|
|
|
|
|
def test_webgpt():
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
|
|
collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
|
|
dataset = WebGPT()
|
|
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
|
|
for batch in dataloader:
|
|
print(batch['input_ids'].shape)
|
|
|
|
|
|
def test_hf_quality():
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
|
|
collate_fn = DataCollatorForSummaryScore(tokenizer, max_length=200)
|
|
dataset = HFSummaryQuality('validation', tokenizer)
|
|
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
|
|
for batch in dataloader:
|
|
print(batch['input_ids'].shape)
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_hf_quality()
|
|
# test_webgpt() |