import json

from datasets import load_dataset

def load_custom_dataset(file_path):
    return load_dataset("json", data_files={"test": file_path},split="test")

def load_jsonlines(file_path):
    with open(file_path, "r") as f:
        return [json.loads(l) for l in f]

def load_json(file_path):
    with open(file_path, "r") as f:
        return json.load(f)

def save_data(data, file_path):
    if file_path.endswith(".jsonl"):
        with open(file_path, 'w', encoding='utf-8') as f:
            for item in data:
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
    elif file_path.endswith('.json'):
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=4)
    else:
        raise ValueError("Please use `.json` or `.jsonl` suffix")

def renames(src_path):
    dataset = load_custom_dataset(src_path)
    dataset = dataset.rename_columns({"outputs": "answer_aliases"})
    dataset.to_json(src_path)