|
import spacy |
|
import zstandard as zstd |
|
import json |
|
import typing |
|
import os |
|
from tqdm import tqdm |
|
import multiprocessing |
|
import random |
|
from langdetect import detect |
|
import argparse |
|
|
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--input_dir', type=str, help='Path to the input file') |
|
args = parser.parse_args() |
|
input_dir = args.input_dir |
|
|
|
|
|
def is_english(text): |
|
try: |
|
lang = detect(text) |
|
return lang == 'en' |
|
except: |
|
return False |
|
|
|
def process_text(texts, model, out_f, lock): |
|
for text in texts: |
|
doc = model(text) |
|
freq_cnt = {} |
|
for e in doc.ents: |
|
if e not in freq_cnt: |
|
freq_cnt[e] = 0 |
|
freq_cnt[e] += 1 |
|
if len(freq_cnt) == 0: |
|
continue |
|
sorted_freq = sorted(freq_cnt.items(), key = lambda x:[1]) |
|
most_freq = sorted_freq[-1][0] |
|
data = {'text':text, 'main_entity':most_freq.text, 'label': most_freq.label_, 'id': most_freq.kb_id_} |
|
json_data = json.dumps(data) |
|
with lock: |
|
out_f.write(json_data + '\n') |
|
out_f.flush() |
|
|
|
def run_ner_linking(texts: typing.List[str], ner_model_path: str): |
|
nlp = spacy.load(ner_model_path) |
|
out_f = open('result/temp_store_data.json', 'w', encoding='utf-8') |
|
lock = multiprocessing.Lock() |
|
processes = [] |
|
|
|
for i in tqdm(range(0, len(texts), 1000)): |
|
p = multiprocessing.Process(target=process_text, args=(texts[i:i+1000], nlp, out_f, lock)) |
|
processes.append(p) |
|
p.start() |
|
|
|
for p in processes: |
|
p.join() |
|
|
|
out_f.close() |
|
return |
|
|
|
wikipedia_out_path='result/wikipedia.json' |
|
subdirectories = [f.path for f in os.scandir(input_dir) if f.is_dir()] |
|
wikipedia_data = [] |
|
for sub_dir in subdirectories: |
|
chunk_dir = sub_dir+'/' |
|
zst_files = [f for f in os.listdir(chunk_dir) if f.endswith('.zst')] |
|
for file in tqdm(zst_files): |
|
with open(chunk_dir+file, 'rb') as compressed_file: |
|
decompressor = zstd.ZstdDecompressor() |
|
with decompressor.stream_reader(compressed_file) as reader: |
|
decompressed_data = reader.read() |
|
for line in decompressed_data.splitlines(): |
|
data = json.loads(line) |
|
|
|
if data['meta']['redpajama_set_name']=='RedPajamaWikipedia': |
|
if is_english(data['text']): |
|
wikipedia_data.append(data) |
|
|
|
with open(wikipedia_out_path, 'w', encoding='utf-8') as f: |
|
for data in wikipedia_data: |
|
json_data = json.dumps(data) |
|
f.write(json_data+'\n') |
|
|
|
wikipedia_data = [] |
|
ner_model_path = 'kc-ner-model' |
|
with open(wikipedia_out_path, 'r', encoding='utf-8') as f: |
|
for line in tqdm(f): |
|
data = json.loads(line) |
|
wikipedia_data.append(data['text']) |
|
run_ner_linking(wikipedia_data, ner_model_path) |
|
|
|
entity_info_path = 'result/entity_info.json' |
|
with open(entity_info_path, 'r', encoding='utf-8') as f: |
|
entity_info = json.load(f) |
|
all_original_data = [] |
|
|
|
category = {} |
|
all_data = [] |
|
with open('result/temp_store_data.json', 'r', encoding='utf-8') as f: |
|
for line in f: |
|
data = json.loads(line) |
|
all_data.append(data) |
|
if data['label'] not in category: |
|
category[data['label']] = [] |
|
category[data['label']].append(data['main_entity']) |
|
|
|
with open('result/processed_data.json', 'w', encoding='utf-8') as f: |
|
for data in tqdm(all_data): |
|
text = data['text'] |
|
main_entity = [data['main_entity']] |
|
if data['id'] in entity_info: |
|
main_entity.extend(entity_info[data['id']]['aliases']) |
|
if len(category[data['label']]) == 1: |
|
continue |
|
replaced_eneity = random.sample(category[data['label']], 1) |
|
while replaced_eneity[0] in main_entity: |
|
replaced_eneity = random.sample(category[data['label']], 1) |
|
for entity in main_entity: |
|
text = text.replace(entity, replaced_eneity[0]) |
|
data = { |
|
'text':text, |
|
'original_main_entity':main_entity, |
|
'replaced_entity':replaced_eneity[0] |
|
} |
|
json_data = json.dumps(data) |
|
f.write(json_data+'\n') |
|
|
|
|
