66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
import os
|
|
import regex as re
|
|
from ast import literal_eval
|
|
import argparse
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Compute corpus statstistics.')
|
|
parser.add_argument('-i', '--input', type=str, required=True, help='Data directory.')
|
|
parser.add_argument('-o', '--output', type=str, required=True, help='Stats file.')
|
|
args = parser.parse_args()
|
|
|
|
stats = calculate_stats(args.input)
|
|
write_stats(stats, args.output)
|
|
|
|
def write_stats(stats, file):
|
|
with open(file, 'wt', encoding='utf-8') as f:
|
|
for record in stats:
|
|
f.write(record + "\n")
|
|
|
|
def calculate_stats(dir):
|
|
src_dir = os.path.join(dir, 'src')
|
|
tgt_dir = os.path.join(dir, 'tgt')
|
|
gold_dir = os.path.join(dir, 'gold')
|
|
stats = []
|
|
header = "\t".join(['id', 'src_sents', 'src_tokens', 'tgt_sents', 'tgt_tokens', 'alignments', '1to1_alignments'])
|
|
stats.append(header)
|
|
for file in sorted(os.listdir(src_dir)):
|
|
if re.match(r'^\d+$', file):
|
|
src_file = os.path.join(src_dir, file + '.tok')
|
|
tgt_file = os.path.join(tgt_dir, file + '.tok')
|
|
gold_file = os.path.join(gold_dir, file + '.align')
|
|
src_sent_num, src_tok_num = count_sent_and_tok_nums(src_file)
|
|
tgt_sent_num, tgt_tok_num = count_sent_and_tok_nums(tgt_file)
|
|
align_num, one_num = count_alignment_nums(gold_file)
|
|
stats.append("\t".join([file, str(src_sent_num), str(src_tok_num), str(tgt_sent_num), str(tgt_tok_num), str(align_num), str(one_num)]))
|
|
return stats
|
|
|
|
def count_alignment_nums(file):
|
|
align_num, one_num = 0, 0
|
|
with open(file, 'rt', encoding="utf-8") as f:
|
|
for line in f:
|
|
align_num += 1
|
|
fields = [x.strip() for x in line.split(':') if len(x.strip())]
|
|
src_len = len(literal_eval(fields[0]))
|
|
tgt_len = len(literal_eval(fields[1]))
|
|
if src_len + tgt_len == 2:
|
|
one_num += 1
|
|
return align_num, one_num
|
|
|
|
def count_sent_and_tok_nums(file):
|
|
sent_num, tok_num = 0, 0
|
|
with open(file, 'rt', encoding='utf-8') as f:
|
|
for line in f:
|
|
sent_num += 1
|
|
line = line.strip()
|
|
tokens = line.split()
|
|
for token in tokens:
|
|
if re.match(r'^\p{P}+$', token):
|
|
continue
|
|
tok_num += 1
|
|
|
|
return sent_num, tok_num
|
|
|
|
if __name__ == '__main__':
|
|
main()
|