Spaces:
Running
on
T4
Running
on
T4
File size: 1,868 Bytes
0c3992e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import os
import os.path as osp
import pickle
import torch
import json
def read_from_file(file_path):
if '.txt' in file_path:
with open(file_path, 'r') as f:
return f.read()
elif '.json' in file_path:
with open(file_path, 'r') as f:
return json.load(f)
elif '.pkl' in file_path:
with open(file_path, 'rb') as f:
return pickle.load(f)
else:
raise NotImplementedError(f'File type not supported: {file_path}')
def write_to_file(file_path, content):
if '.txt' in file_path:
with open(file_path, 'w') as f:
f.write(content)
elif '.json' in file_path:
with open(file_path, 'w') as f:
json.dump(content, f, indent=4)
elif '.pkl' in file_path:
with open(file_path, 'wb') as f:
pickle.dump(content, f)
else:
raise NotImplementedError(f'File type not supported: {file_path}')
def save_files(save_path, **kwargs):
os.makedirs(save_path, exist_ok=True)
for key, value in kwargs.items():
if isinstance(value, dict):
with open(osp.join(save_path, f'{key}.pkl'), 'wb') as f:
pickle.dump(value, f)
elif isinstance(value, torch.Tensor):
torch.save(value, osp.join(save_path, f'{key}.pt'))
else:
pass
def load_files(save_path):
loaded_dict = {}
for file in os.listdir(save_path):
if os.path.isdir(osp.join(save_path, file)):
continue
if file.endswith('.pkl'):
with open(osp.join(save_path, file), 'rb') as f:
loaded_dict[file.split('.')[0]] = pickle.load(f)
elif file.endswith('.pt'):
loaded_dict[file.split('.')[0]] = torch.load(osp.join(save_path, file))
else:
pass
return loaded_dict
|