utils.py

# coding=utf-8
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import math
import random
import time
from urllib.error import URLError
from urllib.parse import urlencode
from urllib.request import Request, urlopen

import numpy as np
import paddle
from tqdm import tqdm


def set_seed(seed):
 paddle.seed(seed)
 random.seed(seed)
 np.random.seed(seed)


class ASRError(Exception):
 pass


def mandarin_asr_api(api_key, secret_key, audio_file, audio_format="wav"):
 """Mandarin ASR

 Args:
 audio_file (str):
 Audio file of Mandarin with sampling rate 16000.
 audio_format (str):
 The file extension of audio_file, 'wav' by default.

 Please refer to https://github.com/Baidu-AIP/speech-demo for more demos.
 """
 # Configurations.
 TOKEN_URL = "http://aip.baidubce.com/oauth/2.0/token"
 ASR_URL = "http://vop.baidu.com/server_api"
 SCOPE = "audio_voice_assistant_get"
 API_KEY = api_key
 SECRET_KEY = secret_key

 # Fetch tokens from TOKEN_URL.
 post_data = urlencode(
 {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
 ).encode("utf-8")

 request = Request(TOKEN_URL, post_data)
 try:
 result_str = urlopen(request).read()
 except URLError as error:
 print("token http response http code : " + str(error.code))
 result_str = error.read()
 result_str = result_str.decode()

 result = json.loads(result_str)
 if "access_token" in result.keys() and "scope" in result.keys():
 if SCOPE and (SCOPE not in result["scope"].split(" ")):
 raise ASRError("scope is not correct!")
 token = result["access_token"]
 else:
 raise ASRError(
 "MAYBE API_KEY or SECRET_KEY not correct: " + "access_token or scope not found in token response"
 )

 # Fetch results by ASR api.
 with open(audio_file, "rb") as speech_file:
 speech_data = speech_file.read()
 length = len(speech_data)
 if length == 0:
 raise ASRError("file %s length read 0 bytes" % audio_file)
 params_query = urlencode({"cuid": "ASR", "token": token, "dev_pid": 1537})
 headers = {"Content-Type": "audio/%s; rate=16000" % audio_format, "Content-Length": length}

 url = ASR_URL + "?" + params_query
 request = Request(url, speech_data, headers)
 try:
 begin = time.time()
 result_str = urlopen(request).read()
 print("Request time cost %f" % (time.time() - begin))
 except URLError as error:
 print("asr http response http code : " + str(error.code))
 result_str = error.read()
 result_str = str(result_str, "utf-8")
 result = json.loads(result_str)

 return result["result"][0]


@paddle.no_grad()
def evaluate(model, metric, data_loader):
 """
 Given a dataset, it evals model and computes the metric.
 Args:
 model(obj:`paddle.nn.Layer`): A model to classify texts.
 metric(obj:`paddle.metric.Metric`): The evaluation metric.
 data_loader(obj:`paddle.io.DataLoader`): The dataset loader which generates batches.
 """
 model.eval()
 metric.reset()
 for batch in data_loader:
 input_ids, token_type_ids, att_mask, pos_ids, start_ids, end_ids = batch
 start_prob, end_prob = model(input_ids, token_type_ids, att_mask, pos_ids)
 start_ids = paddle.cast(start_ids, "float32")
 end_ids = paddle.cast(end_ids, "float32")
 num_correct, num_infer, num_label = metric.compute(start_prob, end_prob, start_ids, end_ids)
 metric.update(num_correct, num_infer, num_label)
 precision, recall, f1 = metric.accumulate()
 model.train()
 return precision, recall, f1


def convert_example(example, tokenizer, max_seq_len):
 """
 example: {
 title
 prompt
 content
 result_list
 }
 """
 encoded_inputs = tokenizer(
 text=[example["prompt"]],
 text_pair=[example["content"]],
 stride=len(example["prompt"]),
 truncation=True,
 max_seq_len=max_seq_len,
 pad_to_max_seq_len=True,
 return_attention_mask=True,
 return_position_ids=True,
 return_dict=False,
 )
 encoded_inputs = encoded_inputs[0]
 offset_mapping = [list(x) for x in encoded_inputs["offset_mapping"]]
 bias = 0
 for index in range(len(offset_mapping)):
 if index == 0:
 continue
 mapping = offset_mapping[index]
 if mapping[0] == 0 and mapping[1] == 0 and bias == 0:
 bias = index
 if mapping[0] == 0 and mapping[1] == 0:
 continue
 offset_mapping[index][0] += bias
 offset_mapping[index][1] += bias
 start_ids = [0 for x in range(max_seq_len)]
 end_ids = [0 for x in range(max_seq_len)]
 for item in example["result_list"]:
 start = map_offset(item["start"] + bias, offset_mapping)
 end = map_offset(item["end"] - 1 + bias, offset_mapping)
 start_ids[start] = 1.0
 end_ids[end] = 1.0

 tokenized_output = [
 encoded_inputs["input_ids"],
 encoded_inputs["token_type_ids"],
 encoded_inputs["position_ids"],
 encoded_inputs["attention_mask"],
 start_ids,
 end_ids,
 ]
 tokenized_output = [np.array(x, dtype="int64") for x in tokenized_output]
 return tuple(tokenized_output)


def map_offset(ori_offset, offset_mapping):
 """
 map ori offset to token offset
 """
 for index, span in enumerate(offset_mapping):
 if span[0] <= ori_offset < span[1]:
 return index
 return -1


def reader(data_path, max_seq_len=512):
 """
 read json
 """
 with open(data_path, "r", encoding="utf-8") as f:
 for line in f:
 json_line = json.loads(line)
 content = json_line["content"]
 prompt = json_line["prompt"]
 # Model Input is aslike: [CLS] Prompt [SEP] Content [SEP]
 # It include three summary tokens.
 if max_seq_len <= len(prompt) + 3:
 raise ValueError("The value of max_seq_len is too small, please set a larger value")
 max_content_len = max_seq_len - len(prompt) - 3
 if len(content) <= max_content_len:
 yield json_line
 else:
 result_list = json_line["result_list"]
 json_lines = []
 accumulate = 0
 while True:
 cur_result_list = []

 for result in result_list:
 if result["start"] + 1 <= max_content_len < result["end"]:
 max_content_len = result["start"]
 break

 cur_content = content[:max_content_len]
 res_content = content[max_content_len:]

 while True:
 if len(result_list) == 0:
 break
 elif result_list[0]["end"] <= max_content_len:
 if result_list[0]["end"] > 0:
 cur_result = result_list.pop(0)
 cur_result_list.append(cur_result)
 else:
 cur_result_list = [result for result in result_list]
 break
 else:
 break

 json_line = {"content": cur_content, "result_list": cur_result_list, "prompt": prompt}
 json_lines.append(json_line)

 for result in result_list:
 if result["end"] <= 0:
 break
 result["start"] -= max_content_len
 result["end"] -= max_content_len
 accumulate += max_content_len
 max_content_len = max_seq_len - len(prompt) - 3
 if len(res_content) == 0:
 break
 elif len(res_content) < max_content_len:
 json_line = {"content": res_content, "result_list": result_list, "prompt": prompt}
 json_lines.append(json_line)
 break
 else:
 content = res_content

 for json_line in json_lines:
 yield json_line


def add_negative_example(examples, texts, prompts, label_set, negative_ratio):
 with tqdm(total=len(prompts)) as pbar:
 for i, prompt in enumerate(prompts):
 negtive_sample = []
 redundants_list = list(set(label_set) ^ set(prompt))
 redundants_list.sort()

 if len(examples[i]) == 0:
 continue
 else:
 actual_ratio = math.ceil(len(redundants_list) / len(examples[i]))

 if actual_ratio <= negative_ratio:
 idxs = [k for k in range(len(redundants_list))]
 else:
 idxs = random.sample(range(0, len(redundants_list)), negative_ratio * len(examples[i]))

 for idx in idxs:
 negtive_result = {"content": texts[i], "result_list": [], "prompt": redundants_list[idx]}
 negtive_sample.append(negtive_result)
 examples[i].extend(negtive_sample)
 pbar.update(1)
 return examples


def construct_relation_prompt_set(entity_name_set, predicate_set):
 relation_prompt_set = set()
 for entity_name in entity_name_set:
 for predicate in predicate_set:
 # The relation prompt is constructed as follows:
 # subject + "的" + predicate
 relation_prompt = entity_name + "的" + predicate
 relation_prompt_set.add(relation_prompt)
 return sorted(list(relation_prompt_set))


def convert_ext_examples(raw_examples, negative_ratio):
 texts = []
 entity_examples = []
 relation_examples = []
 entity_prompts = []
 relation_prompts = []
 entity_label_set = []
 entity_name_set = []
 predicate_set = []

 print("Converting doccano data...")
 with tqdm(total=len(raw_examples)) as pbar:
 for line in raw_examples:
 items = json.loads(line)
 entity_id = 0
 if "data" in items.keys():
 text = items["data"]
 entities = []
 for item in items["label"]:
 entity = {"id": entity_id, "start_offset": item[0], "end_offset": item[1], "label": item[2]}
 entities.append(entity)
 entity_id += 1
 relations = []
 else:
 text, relations, entities = items["text"], items["relations"], items["entities"]
 texts.append(text)

 entity_example = []
 entity_prompt = []
 entity_example_map = {}
 entity_map = {} # id to entity name
 for entity in entities:
 entity_name = text[entity["start_offset"] : entity["end_offset"]]
 entity_map[entity["id"]] = {
 "name": entity_name,
 "start": entity["start_offset"],
 "end": entity["end_offset"],
 }

 entity_label = entity["label"]
 result = {"text": entity_name, "start": entity["start_offset"], "end": entity["end_offset"]}
 if entity_label not in entity_example_map.keys():
 entity_example_map[entity_label] = {
 "content": text,
 "result_list": [result],
 "prompt": entity_label,
 }
 else:
 entity_example_map[entity_label]["result_list"].append(result)

 if entity_label not in entity_label_set:
 entity_label_set.append(entity_label)
 if entity_name not in entity_name_set:
 entity_name_set.append(entity_name)
 entity_prompt.append(entity_label)

 for v in entity_example_map.values():
 entity_example.append(v)

 entity_examples.append(entity_example)
 entity_prompts.append(entity_prompt)

 relation_example = []
 relation_prompt = []
 relation_example_map = {}
 for relation in relations:
 predicate = relation["type"]
 subject_id = relation["from_id"]
 object_id = relation["to_id"]
 # The relation prompt is constructed as follows:
 # subject + "的" + predicate
 prompt = entity_map[subject_id]["name"] + "的" + predicate
 result = {
 "text": entity_map[object_id]["name"],
 "start": entity_map[object_id]["start"],
 "end": entity_map[object_id]["end"],
 }
 if prompt not in relation_example_map.keys():
 relation_example_map[prompt] = {"content": text, "result_list": [result], "prompt": prompt}
 else:
 relation_example_map[prompt]["result_list"].append(result)

 if predicate not in predicate_set:
 predicate_set.append(predicate)
 relation_prompt.append(prompt)

 for v in relation_example_map.values():
 relation_example.append(v)

 relation_examples.append(relation_example)
 relation_prompts.append(relation_prompt)
 pbar.update(1)

 print("Adding negative samples for first stage prompt...")
 entity_examples = add_negative_example(entity_examples, texts, entity_prompts, entity_label_set, negative_ratio)
 if len(predicate_set) != 0:
 print("Constructing relation prompts...")
 relation_prompt_set = construct_relation_prompt_set(entity_name_set, predicate_set)

 print("Adding negative samples for second stage prompt...")
 relation_examples = add_negative_example(
 relation_examples, texts, relation_prompts, relation_prompt_set, negative_ratio
 )
 return entity_examples, relation_examples


def create_dataloader(dataset, mode="train", batch_size=1, batchify_fn=None, trans_fn=None):
 if trans_fn:
 dataset = dataset.map(trans_fn)

 shuffle = True if mode == "train" else False
 if mode == "train":
 batch_sampler = paddle.io.DistributedBatchSampler(dataset, batch_size=batch_size, shuffle=shuffle)
 else:
 batch_sampler = paddle.io.BatchSampler(dataset, batch_size=batch_size, shuffle=shuffle)

 return paddle.io.DataLoader(dataset=dataset, batch_sampler=batch_sampler, collate_fn=batchify_fn, return_list=True)