84 lines
2.2 KiB
Python
84 lines
2.2 KiB
Python
"""
|
|
Normalize extracted text before indexing.
|
|
|
|
Usage:
|
|
python scripts/normalize_text.py <input_dir> <output_dir>
|
|
|
|
Reads JSON files produced by parse_wikipedia.py / parse_pdf.py,
|
|
cleans the text field, writes cleaned files to output_dir.
|
|
"""
|
|
import sys
|
|
import re
|
|
import json
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
|
|
def normalize(text: str) -> str:
|
|
# Collapse unicode whitespace and control characters
|
|
text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f]', '', text)
|
|
|
|
# Normalize various dash/hyphen variants to a regular hyphen
|
|
text = re.sub(r'[—–‒―]', '-', text)
|
|
|
|
# Remove repeated punctuation (e.g. "!!!" → "!")
|
|
text = re.sub(r'([!?.]){2,}', r'\1', text)
|
|
|
|
# Collapse multiple spaces / tabs to a single space
|
|
text = re.sub(r'[ \t]+', ' ', text)
|
|
|
|
# Collapse more than two consecutive newlines
|
|
text = re.sub(r'\n{3,}', '\n\n', text)
|
|
|
|
# Strip leading/trailing whitespace per line
|
|
lines = [line.strip() for line in text.splitlines()]
|
|
text = '\n'.join(lines)
|
|
|
|
return text.strip()
|
|
|
|
|
|
def process_directory(input_dir: str, output_dir: str) -> int:
|
|
out_path = Path(output_dir)
|
|
out_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
count = 0
|
|
for src in sorted(Path(input_dir).glob('*.json')):
|
|
try:
|
|
with open(src, encoding='utf-8') as f:
|
|
doc = json.load(f)
|
|
except Exception as e:
|
|
print(f'Skip {src.name}: {e}')
|
|
continue
|
|
|
|
if 'text' not in doc:
|
|
continue
|
|
|
|
doc['text'] = normalize(doc['text'])
|
|
|
|
if len(doc['text'].split()) < 30:
|
|
continue
|
|
|
|
dest = out_path / src.name
|
|
with open(dest, 'w', encoding='utf-8') as f:
|
|
json.dump(doc, f, ensure_ascii=False, indent=2)
|
|
|
|
count += 1
|
|
if count % 1000 == 0:
|
|
print(f'Normalized {count} documents')
|
|
|
|
return count
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('input_dir')
|
|
parser.add_argument('output_dir')
|
|
args = parser.parse_args()
|
|
|
|
count = process_directory(args.input_dir, args.output_dir)
|
|
print(f'Done: {count} documents written to {args.output_dir}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|