import codecs
import json


def load_data(data_path, data_type="train"):
    """读取数据集
    Args:
        data_path (str): 数据存放路径
        data_type (str, optional): 数据类型. Defaults to "train".
    Returns:
        (json): train和valid中一条数据格式：{"text":"","entity_list":[(start, end, label), (start, end, label)...]}
    """
    if data_type == "train" or data_type == "valid":
        datas = []
        with codecs.open(data_path, encoding="utf-8") as f:
            for line in f:
                line = json.loads(line)
                item = {"text": line["text"], "entity_list": []}
                for k, v in line['label'].items():
                    for spans in v.values():
                        for start, end in spans:
                            item["entity_list"].append((start, end, k))
                datas.append(item)
        return datas
    else:
        return json.load(open(data_path, encoding="utf-8"))