Updated table preprocessing - fixed multiple table issue

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

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)