From be11899de199ebb7113335a9be4fc9a0312ba4a4 Mon Sep 17 00:00:00 2001 From: Katon Minhas Date: Fri, 7 Jun 2024 13:04:05 -0700 Subject: [PATCH] Updated table preprocessing - fixed multiple table issue --- src/table_funcs.py | 32 ++++++++----------- src/test.py | 78 +++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 83 insertions(+), 27 deletions(-) diff --git a/src/table_funcs.py b/src/table_funcs.py index 0d62314..a438cac 100644 --- a/src/table_funcs.py +++ b/src/table_funcs.py @@ -22,7 +22,6 @@ def extract_all_tables(input_dict): tables_dict[filename] = extract_tables(filetext) return tables_dict - def convert_to_dict(table_text): table_text = table_text.strip('{}') table_text_list = table_text.split(']') @@ -32,15 +31,9 @@ def convert_to_dict(table_text): num_elements = 0 for key_value_text in table_text_list: key = key_value_text.split(':')[0].strip(', \'') - value = ':'.join(key_value_text.split(':')[1:]).strip(' [] \'') + value = ':'.join(key_value_text.split(':')[1:]).strip(' \'')+'\']' - f = StringIO(value) - reader = csv.reader(f, delimiter=',') - - value_list = [] - for row in reader: - for r in row: - value_list.append(r) + value_list = ast.literal_eval(value) final_dict[key] = value_list num_elements = len(value_list) @@ -62,6 +55,7 @@ def format_table(table_json, table_size): def align_and_format_tables(text_dict): aligned_text_dict = {} for key, text in text_dict.items(): + aligned_text = text if 'Table Start' in text: table_texts = re.findall(r'-------Table Start--------(.*?)-------Table End--------', text, re.DOTALL) for table_text in table_texts: @@ -75,17 +69,17 @@ def align_and_format_tables(text_dict): table, table_size = convert_to_dict(table_only) table_formatted = format_table(table, table_size) - - # Align - if text.count(pretable) == 2: - # Remove original table - aligned_text = text.replace(table_text, "") - # Align new table - aligned_text = aligned_text.replace(pretable, pretable + table_formatted) - else: - aligned_text = text.replace(table_text, pretable + table_formatted) except: - aligned_text = text + table_formatted = table_text + + # Align + if text.count(pretable) == 2: # One table + # Remove original table + aligned_text = text.replace(table_text, "") + # Align new table + aligned_text = aligned_text.replace(pretable, pretable + table_formatted) + else: + aligned_text = aligned_text.replace(table_text, pretable + table_formatted) aligned_text_dict[key] = aligned_text.replace("-------Table Start--------", "").replace("-------Table End--------", "") else: aligned_text_dict[key] = text diff --git a/src/test.py b/src/test.py index c37fae6..6288640 100644 --- a/src/test.py +++ b/src/test.py @@ -4,6 +4,7 @@ import json import csv from io import StringIO import config +import ast import utils import preprocess @@ -42,18 +43,79 @@ filename = list(input_dict.keys())[0] contract_text = input_dict[filename] print(filename) +def convert_to_dict(table_text): + table_text = table_text.strip('{}') + table_text_list = table_text.split(']') + table_text_list = [item for item in table_text_list if len(item) > 0] + + final_dict = {} + num_elements = 0 + for key_value_text in table_text_list: + key = key_value_text.split(':')[0].strip(', \'') + value = ':'.join(key_value_text.split(':')[1:]).strip(' \'')+'\']' + + value_list = ast.literal_eval(value) + + final_dict[key] = value_list + num_elements = len(value_list) + return final_dict, num_elements + + +def format_table(table_json, table_size): + table_text = "" + keys = [key for key in table_json.keys()] + table_text += ' : '.join(keys) + "\n" + for i in range(table_size): + for key in keys: + if table_json[key][i]: + #table_text += key + ': ' + table_json[key][i] + ', ' + table_text += table_json[key][i] + ' : ' + table_text += r'\n' + return table_text + +def align_and_format_tables(text_dict): + aligned_text_dict = {} + for key, text in text_dict.items(): + aligned_text = text + if 'Table Start' in text: + table_texts = re.findall(r'-------Table Start--------(.*?)-------Table End--------', text, re.DOTALL) + for table_text in table_texts: + try: + # Extract pretable text + pretable = table_text.split('{')[0].strip() + + # Extract and format table text + #table_only = "{" + table_text.split('{', 1)[1].rsplit('}', 1)[0].replace("'", '"') + "}" + table_only = "{" + table_text.split('{', 1)[1].rsplit('}', 1)[0] + "}" + + table, table_size = convert_to_dict(table_only) + table_formatted = format_table(table, table_size) + except: + table_formatted = table_text + + # Align + if text.count(pretable) == 2: # One table + # Remove original table + aligned_text = text.replace(table_text, "") + # Align new table + aligned_text = aligned_text.replace(pretable, pretable + table_formatted) + else: + aligned_text = aligned_text.replace(table_text, pretable + table_formatted) + aligned_text_dict[key] = aligned_text.replace("-------Table Start--------", "").replace("-------Table End--------", "") + else: + aligned_text_dict[key] = text + + return aligned_text_dict + + contract_text = preprocess.clean_newlines(contract_text) text_dict = preprocess.split_text(contract_text) -text_dict = table_funcs.align_and_format_tables(text_dict) +text_dict = align_and_format_tables({'25' : text_dict['25']}) text_dict = preprocess.highlight_rates(text_dict) -print(text_dict['2']) - -# prompt = BOTTOM_UP_PRIMARY(text_dict['2'], '') - -# answer = claude_funcs.invoke_claude_3(prompt, max_tokens=4000) - -# print(answer) +prompt = BOTTOM_UP_PRIMARY(text_dict['25'], '') +answer = claude_funcs.invoke_claude_3(prompt, max_tokens=4000) +print(answer)