54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
|
import pickle
|
||
|
|
from datasets import load_dataset
|
||
|
|
from transformers import AutoTokenizer
|
||
|
|
from tqdm import tqdm
|
||
|
|
from sklearn.preprocessing import MinMaxScaler
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
dataset = load_dataset('csv', sep='\t', data_files={'train': ['./train_huggingface_format.csv'], 'test': ['./dev-0_huggingface_format.csv']})
|
||
|
|
test_dataset_A = load_dataset('csv', sep='\t', data_files='test-A_huggingface_format.csv')
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained('without_date/checkpoint-395000')
|
||
|
|
|
||
|
|
def tokenize_function(examples):
|
||
|
|
t = tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
|
||
|
|
return t
|
||
|
|
|
||
|
|
test_tokenized_datasets_A = test_dataset_A.map(tokenize_function, batched=True)
|
||
|
|
tokenized_datasets = dataset.map(tokenize_function, batched=True)
|
||
|
|
|
||
|
|
train_dataset = tokenized_datasets["train"].shuffle(seed=42)
|
||
|
|
eval_dataset_full = tokenized_datasets["test"]
|
||
|
|
eval_dataset_small = tokenized_datasets["test"].select(range(2000))
|
||
|
|
test_dataset_A = test_tokenized_datasets_A["train"]
|
||
|
|
|
||
|
|
|
||
|
|
scalers = dict()
|
||
|
|
scalers['year'] = MinMaxScaler().fit(np.array(train_dataset['year']).reshape(-1,1))
|
||
|
|
|
||
|
|
def add_scaled(example):
|
||
|
|
for factor in ('year',):
|
||
|
|
example[factor + '_scaled'] = scalers[factor].transform(np.array(example[factor]).reshape(-1,1)).reshape(1,-1)[0].item()
|
||
|
|
return example
|
||
|
|
|
||
|
|
train_dataset = train_dataset.map(add_scaled)
|
||
|
|
eval_dataset_full = eval_dataset_full.map(add_scaled)
|
||
|
|
eval_dataset_small = eval_dataset_small.map(add_scaled)
|
||
|
|
#test_dataset_A = test_dataset_A.map(add_scaled)
|
||
|
|
|
||
|
|
|
||
|
|
with open('train_dataset.pickle','wb') as f_p:
|
||
|
|
pickle.dump(train_dataset, f_p)
|
||
|
|
|
||
|
|
with open('eval_dataset_small.pickle','wb') as f_p:
|
||
|
|
pickle.dump(eval_dataset_small, f_p)
|
||
|
|
|
||
|
|
with open('eval_dataset_full.pickle','wb') as f_p:
|
||
|
|
pickle.dump(eval_dataset_full, f_p)
|
||
|
|
|
||
|
|
with open('test_dataset_A.pickle','wb') as f_p:
|
||
|
|
pickle.dump(test_dataset_A, f_p)
|
||
|
|
|
||
|
|
with open('scalers.pickle','wb') as f_p:
|
||
|
|
pickle.dump(scalers, f_p)
|