cjvt-srl-tagging/tools/gen_json.py

389 lines
13 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import pickle
from pathlib import Path
from parser.parser import Parser
import configparser
import json
import sys
import logging
from multiprocessing import Pool
# parse config
config = configparser.ConfigParser()
config.read("tools.cfg.ssj500k2.3")
ORIGPATH = Path(config["tools"]["ssj500k_orig_folder"])
INPATH = Path(config["tools"]["ssj500k_srl"])
OUTPATH = Path(config["tools"]["ssj500k_json"])
INTERNAL_DATA = Path(config["tools"]["internal_data"])
DEBUG = config["tools"]["debug"] == "True"
CPU_CORES = int(config["tools"]["cpu_cores"])
LOGFILE = Path(config["tools"]["logfile"]).absolute()
LOGFILE.touch(exist_ok=True)
LOGFILE.resolve()
logging.basicConfig(filename=str(LOGFILE), level=logging.INFO)
def get_origfile(filename):
for origfile in ORIGPATH.iterdir():
if filename.name.split('.')[0] == origfile.name.split('.')[0]:
return origfile
raise FileNotFoundError
def extract_sentences(line_reader):
acc = []
# last char in line is \n, remove it
for line in [x.decode("utf-8")[:-1].split('\t') for x in line_reader]:
if len(line) == 1: # empty line
tmp = acc
acc = []
yield tmp
else:
acc.append(line)
def to_sentence(sentence_arr):
return " ".join([token[1] for token in sentence_arr])
def match_sentence_id(sentence, orig_dict):
for k, e in orig_dict.items():
orig_sentence = " ".join(token[2] for token in e["tokens"])
if sentence == orig_sentence:
return k
raise KeyError
def match_sentence_id_giga(sentence, orig_dict):
for k, e in orig_dict.items():
# orig_sentence = " ".join(token[2] for token in e["tokens"])
if sentence == e["text"]:
return k
raise KeyError
def get_dep_rel(token):
logging.debug(token)
for i, field in enumerate(token[14:]):
if field != "_":
return {
"arg": field,
"from": i, # i-th predicate in sentence
"dep": token[0],
}
return None
def handle_file_old(infile_tpl):
i = infile_tpl[0]
infile = infile_tpl[1]
outfile = (OUTPATH / infile.name).with_suffix(".json")
origfile = get_origfile(infile)
orig_dict = par.parse_tei(origfile)
with infile.open("rb") as fp:
outdata = {}
for sentence_arr in extract_sentences(fp.readlines()):
# tsv dropped sentence ids, match the ID, using original data
sid = match_sentence_id(to_sentence(sentence_arr), orig_dict)
outdata[sid] = []
# find all predicate indices in the sentence
predicates = []
for token in sentence_arr:
if token[12] == "Y":
predicates += [token[0]] # idx
deprel = get_dep_rel(token)
if deprel is not None:
outdata[sid].append(deprel)
# deprel["from"] points to n-th predicate
# replace with predicate's token index
for deprel in outdata[sid]:
deprel["from"] = predicates[deprel["from"]]
if DEBUG:
print(to_sentence(sentence_arr))
print(outdata[sid])
print(sid)
print()
print()
with outfile.open("w") as fp:
json.dump(outdata, fp)
logging.info("SRL relations written to: {}".format(outfile))
def handle_file(whole_input):
# sentence_id = whole_input[0][3]
# orig_infile = whole_input[0][1]
sentence_id = whole_input[3]
orig_infile = whole_input[1]
# origfile = origfiles[0][1]
# infile_tpl = infile_tpl[0]
# i = infile_tpl[0]
# infile = infile_tpl[1]
outfile = (OUTPATH / orig_infile.name).with_suffix(".json")
if outfile.exists():
return
# origfile = get_origfile()
orig_dict = par.parse_tei(orig_infile)
outdata = {}
gen = srl_multiple_files_sentences_generator(sentence_id)
# gen = srl_multiple_files_sentences_generator(whole_input[1])
mismatch_sentences = 0
for sentence_i, (orig_id, orig_val) in enumerate(orig_dict.items()):
if orig_id == 'GF0014802.2685.7':
print('PAUSE')
# look at neighbouring sentences if they are correct
sentence, sentence_arr = next(gen)
# orig_sentence = " ".join(token[2] for token in e["tokens"])
assert sentence.replace(' ', '') == orig_val['text']
# if i != 10 and i != 0:
# print('OK!')
sid = orig_id
outdata[sid] = []
# find all predicate indices in the sentence
predicates = []
for token in sentence_arr:
if token[12] == "Y":
predicates += [token[0]] # idx
deprel = get_dep_rel(token)
if deprel is not None:
outdata[sid].append(deprel)
# deprel["from"] points to n-th predicate
# replace with predicate's token index
for deprel in outdata[sid]:
deprel["from"] = predicates[deprel["from"]]
if DEBUG:
print(to_sentence(sentence_arr))
print(outdata[sid])
print(sid)
print()
print()
if mismatch_sentences > 0:
if mismatch_sentences / len(orig_dict.items()) < 0.1:
print('Slight mismatch - %d' % sentence_id)
print(whole_input)
print('ABS mitigated %d' % mismatch_sentences)
print('------------------------------------------------')
else:
print('ERRRRRRRRRRRRRRRROOOOOOORRRRRRRRRRR')
print('Big mismatch - %d' % sentence_id)
print(whole_input)
print('ABS mitigated errors:')
print(mismatch_sentences)
print('------------------------------------------------')
with outfile.open("w") as fp:
json.dump(outdata, fp)
logging.info("SRL relations written to: {}".format(outfile))
def count_orig_file_sentences(filename):
if os.path.exists(os.path.join(INTERNAL_DATA, 'orig_chunks', filename[1].name)):
return
print(filename[0])
orig_dict = par.parse_tei(filename[1])
# return filename[0], filename[1], len(orig_dict)
with open(os.path.join(INTERNAL_DATA, 'orig_chunks', filename[1].name), 'wb') as output:
pickle.dump((filename[0], filename[1], len(orig_dict)), output)
def count_srl_file_sentences(filename):
if os.path.exists(os.path.join(INTERNAL_DATA, 'srl_chunks', filename[1].name)):
return
print(filename[0])
num_sentences = 0
with filename[1].open("r") as fp:
for line in fp:
if line == '\n':
num_sentences += 1
# return filename[0], filename[1], num_sentences
with open(os.path.join(INTERNAL_DATA, 'srl_chunks', filename[1].name), 'wb') as output:
pickle.dump((filename[0], filename[1], num_sentences), output)
def srl_sentences_generator(infile, curr_index, sen_start_index):
with infile.open("rb") as fp:
outdata = {}
for sentence_arr in extract_sentences(fp.readlines()):
if curr_index < sen_start_index:
curr_index += 1
else:
yield to_sentence(sentence_arr), sentence_arr
yield None
def srl_multiple_files_sentences_generator(sentence_id): # srl_files):
sentence_id = max(0, sentence_id - 10)
for i, srl_file in enumerate(srl_file_sizes):
if sentence_id >= srl_file[3] and sentence_id < srl_file[3] + srl_file[2]:
srl_files = srl_file_sizes[i:]
break
for file_info in srl_files:
# srl_gen = srl_sentences_generator(file_info[1], file_info[3], file_info[4])
srl_gen = srl_sentences_generator(file_info[1], file_info[3], sentence_id)
el = next(srl_gen)
while el is not None:
yield el
el = next(srl_gen)
yield None
# main
par = Parser()
OUTPATH.mkdir(exist_ok=True)
infiles = list(enumerate([x for x in iter(sorted(INPATH.iterdir())) if x.is_file()]))
logging.info("Generating JSON SRL files from {} tsv files.".format(len(infiles)))
origfiles = []
for subdir, dirs, files in os.walk(ORIGPATH):
for file in files:
origfiles.append(Path(os.path.join(subdir, file)))
origfiles=list(enumerate(sorted(origfiles)))
##### REMOVE ############
# origfiles = origfiles[:3]
# count sentences in orig (if not counted before)
# os.remove(os.path.join(INTERNAL_DATA, 'orig_counted_sentences.pkl'))
if not os.path.exists(os.path.join(INTERNAL_DATA, 'orig_counted_sentences.pkl')):
# srl_file_sizes = {}
if not os.path.exists(os.path.join(INTERNAL_DATA, 'orig_chunks')):
os.makedirs(os.path.join(INTERNAL_DATA, 'orig_chunks'))
# with Pool(CPU_CORES) as p:
# # p.map(handle_file, infiles)
# p.map(count_orig_file_sentences, origfiles)
for i in range(len(origfiles)):
count_orig_file_sentences(origfiles[i])
orig_file_sizes = []
for x in iter(sorted(Path(os.path.join(INTERNAL_DATA, 'orig_chunks')).iterdir())):
print(x.name)
if x.is_file():
with x.open('rb') as pkl_small_file:
orig_file_sizes.append(pickle.load(pkl_small_file))
# orig_file_sizes = list(enumerate([x for x in iter(sorted(INPATH.iterdir())) if x.is_file()]))
print("Sorting orig files")
orig_file_sizes = sorted(orig_file_sizes)
total_size = 0
orig_file_sizes_final = []
print("Calculating orig files size")
for n, pa, si in orig_file_sizes:
orig_file_sizes_final.append((n, pa, si, total_size))
total_size += si
orig_file_sizes = orig_file_sizes_final
print("Saving orig files size")
with open(os.path.join(INTERNAL_DATA, 'orig_counted_sentences.pkl'), 'wb') as output:
pickle.dump(orig_file_sizes, output)
print("Orig files saved")
else:
with open(os.path.join(INTERNAL_DATA, 'orig_counted_sentences.pkl'), 'rb') as pkl_file:
orig_file_sizes = pickle.load(pkl_file)
# count sentences in srl (if not counted before)
# os.remove(os.path.join(INTERNAL_DATA, 'srl_counted_sentences.pkl'))
if not os.path.exists(os.path.join(INTERNAL_DATA, 'srl_counted_sentences.pkl')):
# srl_file_sizes = {}
if not os.path.exists(os.path.join(INTERNAL_DATA, 'srl_chunks')):
os.makedirs(os.path.join(INTERNAL_DATA, 'srl_chunks'))
# with Pool(CPU_CORES) as p:
# # p.map(handle_file, infiles)
# p.map(count_srl_file_sentences, infiles)
for i in range(len(infiles)):
count_srl_file_sentences(infiles[i])
srl_file_sizes = []
for x in iter(sorted(Path(os.path.join(INTERNAL_DATA, 'srl_chunks')).iterdir())):
print(x.name)
if x.is_file():
with x.open('rb') as pkl_small_file:
srl_file_sizes.append(pickle.load(pkl_small_file))
print("Sorting srl files")
srl_file_sizes = sorted(srl_file_sizes)
total_size = 0
srl_file_sizes_final = []
print("Calculating srl files size")
for n, pa, si in srl_file_sizes:
srl_file_sizes_final.append((n, pa, si, total_size))
total_size += si
srl_file_sizes = srl_file_sizes_final
print("Saving srl files size")
with open(os.path.join(INTERNAL_DATA, 'srl_counted_sentences.pkl'), 'wb') as output:
pickle.dump(srl_file_sizes, output)
print("Srl files saved")
else:
with open(os.path.join(INTERNAL_DATA, 'srl_counted_sentences.pkl'), 'rb') as pkl_file:
srl_file_sizes = pickle.load(pkl_file)
# print(len(orig_file_sizes))
# print('asd' + 2)
# inputs = []
# srl_i = 0
# srl_file = srl_file_sizes[srl_i]
# for orig_i, orig_path, orig_size, orig_first_sent_i in orig_file_sizes:
# interesting_srl_files = []
# # beginning of srl chunk in range of orig chunk or ending of srl chunk in range of orig chunk
# # while srl_file[3] >= orig_first_sent_i and srl_file[3] < orig_first_sent_i + orig_size or \
# # srl_file[3] + srl_file[2] - 1 >= orig_first_sent_i and srl_file[3] + srl_file[2] - 1 < orig_first_sent_i + orig_size:
# while srl_file[3] < orig_first_sent_i + orig_size and srl_file[3] + srl_file[2] > orig_first_sent_i:
# # if beginning of file is in
# if srl_file[3] > orig_first_sent_i:
# interesting_srl_files.append((srl_file[0], srl_file[1], srl_file[2], srl_file[3], srl_file[3]))
# # print('if %d' % srl_file[3])
# else:
# interesting_srl_files.append((srl_file[0], srl_file[1], srl_file[2], srl_file[3], orig_first_sent_i))
# # print('else %d' % orig_first_sent_i)
#
# if orig_first_sent_i + orig_size >= srl_file[3] + srl_file[2]:
# srl_i += 1
# if srl_i < len(srl_file_sizes):
# srl_file = srl_file_sizes[srl_i]
# else:
# break
# # print(srl_i)
# # print('a ' + 2)
# else:
# break
#
# inputs.append([[orig_i, orig_path, orig_size, orig_first_sent_i], interesting_srl_files])
# print(inputs[-1])
# srl_gen = srl_sentences_generator(srl_file_sizes[0][1], 0, 533)
# a = next(srl_gen)
# b = next(srl_gen)
# c = next(srl_gen)
print('beginning processing')
with Pool(CPU_CORES) as p:
# p.map(handle_file, inputs)
p.map(handle_file, orig_file_sizes)
# for of in orig_file_sizes:
# handle_file(of)
logging.info("Finished generating .json files.")