88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
|
from utils import read_local
|
||
|
|
import preprocess
|
||
|
|
import preprocessing_funcs
|
||
|
|
import keywords
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
from keywords import KEYWORD_MAPPINGS
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
# import tqdm
|
||
|
|
|
||
|
|
|
||
|
|
def parse_arguments():
|
||
|
|
parser = argparse.ArgumentParser(description="Smart Chunking Tester")
|
||
|
|
parser.add_argument(
|
||
|
|
"--input_dir", help="Input directory (local path or S3 URI)", default="src/ip2"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--output_dir", help="Output directory (local)", default="src/output_chunks"
|
||
|
|
)
|
||
|
|
parser.add_argument("--keyword", help="Keyword to be used for chunking")
|
||
|
|
parser.add_argument(
|
||
|
|
"--case_sensitive",
|
||
|
|
help="case sensitivity",
|
||
|
|
action=argparse.BooleanOptionalAction,
|
||
|
|
)
|
||
|
|
return parser.parse_args()
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
args = parse_arguments()
|
||
|
|
input_dir = args.input_dir
|
||
|
|
output_dir = args.output_dir # src/output_chunks
|
||
|
|
case_sensitive = args.case_sensitive
|
||
|
|
|
||
|
|
if case_sensitive is None: # This shouldn't have to be assigned
|
||
|
|
case_sensitive = False
|
||
|
|
|
||
|
|
keyword = args.keyword
|
||
|
|
keyword = list(map(str, keyword.split(",")))
|
||
|
|
keyword = [kw.strip() for kw in keyword]
|
||
|
|
print("Generating Chunks for keywords --> ", keyword)
|
||
|
|
print("Case sensitive -->", case_sensitive)
|
||
|
|
print(type(case_sensitive))
|
||
|
|
|
||
|
|
kws = KEYWORD_MAPPINGS
|
||
|
|
|
||
|
|
retained_rates = []
|
||
|
|
|
||
|
|
for file in os.listdir(input_dir):
|
||
|
|
full_path = os.path.join(input_dir, file)
|
||
|
|
contract_text = read_local(full_path)
|
||
|
|
contract_text = preprocessing_funcs.clean_newlines(contract_text)
|
||
|
|
contract_text = preprocessing_funcs.clean_law_symbols(contract_text)
|
||
|
|
text_dict = preprocessing_funcs.split_text(
|
||
|
|
contract_text
|
||
|
|
) # return a dictionary with keys - page_num (str), values as the page_text
|
||
|
|
|
||
|
|
test_chunck = preprocessing_funcs.smart_chunk_ac(
|
||
|
|
text_dict=text_dict,
|
||
|
|
keyword_mappings={
|
||
|
|
"place_holder": {
|
||
|
|
"methodology": "or",
|
||
|
|
"keywords": keyword,
|
||
|
|
"case_sensitive": case_sensitive,
|
||
|
|
}
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
retained_rate = len(test_chunck["place_holder"]) / len(contract_text)
|
||
|
|
retained_rates.append(retained_rate)
|
||
|
|
|
||
|
|
print(
|
||
|
|
f"\nfrom file {file} {len(contract_text)} characters were retrieved;\n{len(test_chunck['place_holder'])} were retained by smart chunking on keywords:\n{keyword}"
|
||
|
|
)
|
||
|
|
print(f"Smart chunking reduced the document by {100*(1-retained_rate):0.2f}%")
|
||
|
|
|
||
|
|
with open(f"{output_dir}/{file}_CHUNKED.txt", "w") as f:
|
||
|
|
f.write(test_chunck["place_holder"])
|
||
|
|
|
||
|
|
print("avg reduction:", 1 - (sum(retained_rates) / len(retained_rates)))
|
||
|
|
print("max retention:", max(retained_rates))
|
||
|
|
print("min retention:", min(retained_rates))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|