utils/input_process.py

"""
사용자 입력을 가공하는 모듈
"""
import copy
import re

from torch.utils.data import DataLoader, Dataset


class ISDataset(Dataset):
 """
 Dataset subclass for Identifying speaker.
 """
 def __init__(self, data_list):
 super(ISDataset, self).__init__()
 self.data = data_list

 def __len__(self):
 return len(self.data)

 def __getitem__(self, idx):
 return self.data[idx]


def make_instance_list(text: str, ws=10) -> list:
 """
 입력받은 문장을 기초적인 인스턴스 리스트로 만들어줍니다.
 """
 lines = text.splitlines()
 max_line = len(lines)

 utterance = ['"', '“', '‘']
 instance_num = []

 for idx, line in enumerate(lines):
 if any(u in line for u in utterance):
 instance_num.append(idx)

 instance = [[] for _ in range(len(instance_num))]

 for i, num in enumerate(instance_num):
 if num - ws <= 0 and num + ws + 1 < max_line:
 instance[i] += ([''] * (ws - num))
 instance[i] +=(lines[:num + 1 + ws])
 elif num - ws <= 0 and num + ws + 1 >= max_line:
 instance[i] += ([''] * (ws - num))
 instance[i] +=(lines)
 instance[i] += ([''] * (ws * 2 - len(instance[i]) + 1))
 elif num + ws + 1 >= max_line:
 instance[i] +=(lines[num-ws:max_line+1])
 instance[i] += ([''] * (num + ws + 1 - max_line))
 else:
 instance[i] += (lines[num-ws:num + ws + 1])

 return instance, instance_num


def NML(seg_sents, mention_positions, ws):
 """
 Nearest Mention Location
 """
 def word_dist(pos):
 """
 The word level distance between quote and the mention position
 """
 if pos[0] == ws:
 w_d = ws * 2
 elif pos[0] < ws:
 w_d = sum(len(
 sent) for sent in seg_sents[pos[0] + 1:ws]) + len(seg_sents[pos[0]][pos[1] + 1:])
 else:
 w_d = sum(
 len(sent) for sent in seg_sents[ws + 1:pos[0]]) + len(seg_sents[pos[0]][:pos[1]])
 return w_d

 sorted_positions = sorted(mention_positions, key=lambda x: word_dist(x))

 return sorted_positions[0]


def max_len_cut(seg_sents, mention_pos, max_len):
 sent_char_lens = [sum(len(word) for word in sent) for sent in seg_sents]
 sum_char_len = sum(sent_char_lens)

 running_cut_idx = [len(sent) - 1 for sent in seg_sents]

 while sum_char_len > max_len:
 max_len_sent_idx = max(list(enumerate(sent_char_lens)), key=lambda x: x[1])[0]

 if max_len_sent_idx == mention_pos[0] and running_cut_idx[max_len_sent_idx] == mention_pos[1]:
 running_cut_idx[max_len_sent_idx] -= 1

 if max_len_sent_idx == mention_pos[0] and running_cut_idx[max_len_sent_idx] < mention_pos[1]:
 mention_pos[1] -= 1

 reduced_char_len = len(
 seg_sents[max_len_sent_idx][running_cut_idx[max_len_sent_idx]])
 sent_char_lens[max_len_sent_idx] -= reduced_char_len
 sum_char_len -= reduced_char_len

 del seg_sents[max_len_sent_idx][running_cut_idx[max_len_sent_idx]]

 running_cut_idx[max_len_sent_idx] -= 1

 return seg_sents, mention_pos


def seg_and_mention_location(raw_sents_in_list, alias2id):
 character_mention_poses = {}
 seg_sents = []
 id_pattern = ['&C{:02d}&'.format(i) for i in range(51)]

 for sent_idx, sent in enumerate(raw_sents_in_list):
 raw_sent_with_split = sent.split()

 for word_idx, word in enumerate(raw_sent_with_split):
 match = re.search(r'&C\d{1,2}&', word)

 if match:
 result = match.group(0)

 if alias2id[result] in character_mention_poses:
 character_mention_poses[alias2id[result]].append([sent_idx, word_idx])
 else:
 character_mention_poses[alias2id[result]] = [[sent_idx, word_idx]]

 seg_sents.append(raw_sent_with_split)

 name_list_index = list(character_mention_poses.keys())

 return seg_sents, character_mention_poses, name_list_index


def create_css(seg_sents, candidate_mention_poses, ws=10):
 """
 Create candidate-specific segments for each candidate in an instance.
 """
 # assert len(seg_sents) == ws * 2 + 1

 many_css = []
 many_sent_char_lens = []
 many_mention_poses = []
 many_quote_idxes = []
 many_cut_css = []

 for candidate_idx in candidate_mention_poses.keys():
 nearest_pos = NML(seg_sents, candidate_mention_poses[candidate_idx], ws)

 if nearest_pos[0] <= ws:
 CSS = copy.deepcopy(seg_sents[nearest_pos[0]:ws + 1])
 mention_pos = [0, nearest_pos[1]]
 quote_idx = ws - nearest_pos[0]
 else:
 CSS = copy.deepcopy(seg_sents[ws:nearest_pos[0] + 1])
 mention_pos = [nearest_pos[0] - ws, nearest_pos[1]]
 quote_idx = 0

 cut_CSS, mention_pos = max_len_cut(CSS, mention_pos, 510)
 sent_char_lens = [sum(len(word) for word in sent) for sent in cut_CSS]

 mention_pos_left = sum(sent_char_lens[:mention_pos[0]]) + sum(
 len(x) for x in cut_CSS[mention_pos[0]][:mention_pos[1]])
 mention_pos_right = mention_pos_left + len(cut_CSS[mention_pos[0]][mention_pos[1]])
 mention_pos = (mention_pos[0], mention_pos_left, mention_pos_right, mention_pos[1])
 cat_CSS = ' '.join([' '.join(sent) for sent in cut_CSS])

 many_css.append(cat_CSS)
 many_sent_char_lens.append(sent_char_lens)
 many_mention_poses.append(mention_pos)
 many_quote_idxes.append(quote_idx)
 many_cut_css.append(cut_CSS)

 return many_css, many_sent_char_lens, many_mention_poses, many_quote_idxes, many_cut_css


def input_data_loader(instances: list, alias2id) -> DataLoader:
 """
 나눠진 데이터를 맞추기 위해 가공
 """
 data_list = []

 for instance in instances:
 seg_sents, candidate_mention_poses, name_list_index = seg_and_mention_location(
 instance, alias2id)
 css, sent_char_lens, mention_poses, quote_idxes, cut_css = create_css(
 seg_sents, candidate_mention_poses)

 data_list.append((seg_sents, css, sent_char_lens, mention_poses, quote_idxes,
 cut_css, name_list_index))

 data_loader = DataLoader(ISDataset(data_list), batch_size=1, collate_fn=lambda x: x[0])

 return data_loader


def make_ner_input(text, chunk_size=500) -> list:
 """
 문장을 New Lines 기준으로 나누어 줍니다.
 chunk size보다 문장이 길 경우, 마지막 문장은 뒤에서 chunk size 만큼 추가합니다.
 """
 count_text = chunk_size
 max_text = len(text)
 newline_position = []

 while count_text < max_text:
 sentence = text[:count_text]
 last_newline_position = sentence.rfind('\n')
 newline_position.append(last_newline_position)
 count_text = last_newline_position + chunk_size

 split_sentences = []
 start_num = 0

 for _, num in enumerate(newline_position):
 split_sentences.append(text[start_num:num])
 start_num = num

 if max_text % chunk_size != 0:
 f_sentence = text[max_text-500:]
 first_newline_position = max_text-500 + f_sentence.find('\n')
 split_sentences.append(text[first_newline_position:])

 return split_sentences


def making_script(text, speaker:list, instance_num:list) -> str:
 """
 스크립트를 만드는 함수
 """
 lines = text.splitlines()
 for num, people in zip(instance_num, speaker):
 lines[num] = f'{people}: {lines[num]}'
 return lines