sparrow-parse 0.1.9__py3-none-any.whl → 0.2.0__py3-none-any.whl

sparrow_parse/__init__.py

@@ -1 +1 @@
-__version__ = '0.1.9'
+__version__ = '0.2.0'

sparrow_parse/data/invoice_1_table.txt

@@ -0,0 +1,9 @@
+[
+    '<table><thead><th>No.</th><th>Description</th><th>Qty</th><th>UM</th><th>Net price</th><th>Net worth</th><th>VAT [%]</th><th>Gross worth</th></thead><tr><td></td><td>Wine Glasses Goblets Pair Clear
+Glass</td><td>5,00</td><td>eacn</td><td>12,00</td><td>60,00</td><td>10%</td><td>66,00</td></tr><tr><td></td><td>With Hooks Stemware Storage Multiple Uses Iron Wine Rack Hanging
+Glass</td><td>4,00</td><td>eacn</td><td>28,08</td><td>112,32</td><td>10%</td><td>123,55</td></tr><tr><td></td><td>Replacement Corkscrew Parts Spiral Worm Wine Opener Bottle
+Houdini</td><td>1,00</td><td>eacn</td><td>7,50</td><td>7,50</td><td>10%</td><td>8,25</td></tr><tr><td></td><td>HOME ESSENTIALS GRADIENT STEMLESS WINE GLASSES SET OF 4 20 FL OZ (591 ml)
+NEW</td><td>1,00</td><td>eacn</td><td>12,99</td><td>12,99</td><td>10%</td><td>14,29</td></tr></table>',
+    '<table><thead><th>VAT</th><th>[%]</th><th>Net worth</th><th>VAT</th><th>Gross worth</th></thead><tr><td></td><td>10%</td><td>192,81</td><td>19,28</td><td>212,09</td></tr><tr><td colspan="2">Total</td><td>$ 192,81</td><td>$
+19,28</td><td>$ 212,09</td></tr></table>'
+]

sparrow_parse/extractor/extractor_helper.py

@@ -0,0 +1,368 @@
+from bs4 import BeautifulSoup
+from sentence_transformers import SentenceTransformer, util
+import pandas as pd
+import re
+from io import StringIO
+
+
+def merge_html_table_headers(html_table, column_keywords, debug=False):
+    soup = BeautifulSoup(html_table, 'html.parser')
+
+    # Find all thead elements
+    theads = soup.find_all('thead')
+
+    if len(theads) > 1 and column_keywords is not None:
+        html_table = update_table_header_colspan(html_table)
+        html_table = merge_table_header_thead(html_table)
+        html_table = merge_colspan_columns(html_table)
+        html_table = normalize_html_table(html_table, debug)
+        html_table = fix_rowspan_elements(html_table)
+        html_table = merge_rows_with_rowspan(html_table)
+        html_table = detect_and_remove_junk_columns(html_table, column_keywords, debug)
+    else:
+        # If there is only one thead, return the original table
+        return html_table
+
+    return html_table
+
+
+def update_table_header_colspan(html_table):
+    soup = BeautifulSoup(html_table, 'html.parser')
+    theads = soup.find_all('thead')
+
+    for thead in theads:
+        for th in thead.find_all('th'):
+            colspan = th.get('colspan')
+            if colspan and int(colspan) > 1:
+                colspan_count = int(colspan)
+                th['colspan'] = 1
+                for _ in range(colspan_count - 1):
+                    new_th = soup.new_tag('th')
+                    th.insert_after(new_th)
+
+    return str(soup)
+
+
+def merge_table_header_thead(html_table):
+    soup = BeautifulSoup(html_table, 'html.parser')
+    theads = soup.find_all('thead')
+
+    primary_thead = theads[0]
+    secondary_thead = theads[1]
+
+    primary_ths = primary_thead.find_all('th')
+    secondary_ths = secondary_thead.find_all('th')
+
+    for i, th in enumerate(primary_ths):
+        if i < len(secondary_ths):
+            primary_text = th.text.strip()
+            secondary_text = secondary_ths[i].text.strip()
+            if primary_text and secondary_text:
+                th.string = (primary_text + ' ' + secondary_text).strip()
+            elif not primary_text and secondary_text:
+                th.string = secondary_text
+        # Remove colspan and rowspan attributes
+        th.attrs.pop('colspan', None)
+        th.attrs.pop('rowspan', None)
+
+    secondary_thead.decompose()
+
+    return str(soup)
+
+
+def merge_colspan_columns(html_table):
+    # Parse the HTML
+    soup = BeautifulSoup(html_table, 'html.parser')
+
+    # Process colspan attributes by adding empty <td> elements
+    for row in soup.find_all('tr'):
+        cols = []
+        for cell in row.find_all(['th', 'td']):
+            colspan = int(cell.get('colspan', 1))
+            # Add the cell and additional empty cells if colspan is greater than 1
+            cols.append(cell)
+            for _ in range(colspan - 1):
+                new_td = soup.new_tag('td')
+                cols.append(new_td)
+            # Remove the colspan attribute
+            if cell.has_attr('colspan'):
+                del cell['colspan']
+
+        # Replace the row's children with the updated cells
+        row.clear()
+        row.extend(cols)
+
+    return str(soup)
+
+
+def normalize_html_table(html, debug = False):
+    soup = BeautifulSoup(html, 'html.parser')
+
+    # Find the header row and count the number of cells
+    header = soup.find('thead').find_all(['th', 'td'])
+    header_cell_count = len(header)
+
+    if debug:
+        # Print the number of header cells
+        print(f"Number of cells in header: {header_cell_count}")
+
+    # Find all rows in the table body
+    rows = soup.find_all('tr')
+
+    for row in rows:
+        cells = row.find_all(['td', 'th'])
+        if len(cells) > header_cell_count:
+            extra_cells = len(cells) - header_cell_count
+            for cell in cells:
+                if cell.text.strip() == '' and extra_cells > 0:
+                    cell.decompose()
+                    extra_cells -= 1
+        elif len(cells) < header_cell_count:
+            missing_cells = header_cell_count - len(cells)
+            for _ in range(missing_cells):
+                new_cell = soup.new_tag('td')
+                row.insert(0, new_cell)
+
+    return str(soup)
+
+
+def fix_rowspan_elements(html_table):
+    # Parse the HTML table
+    soup = BeautifulSoup(html_table, 'html.parser')
+
+    # Find all table rows
+    rows = soup.find_all('tr')
+
+    # Dictionary to store rows with rowspan elements
+    rowspan_dict = {}
+
+    # Iterate over each row
+    for row_index, row in enumerate(rows):
+        # Find all cells in the row
+        cells = row.find_all(['td', 'th'])
+
+        # Iterate over each cell
+        for cell_index, cell in enumerate(cells):
+            # Check if the cell has a rowspan attribute
+            if cell.has_attr('rowspan'):
+                # Store the rowspan value and cell position
+                rowspan_value = int(cell['rowspan'])
+                if row_index not in rowspan_dict:
+                    rowspan_dict[row_index] = []
+                rowspan_dict[row_index].append((cell_index, rowspan_value))
+
+    # List to store the number of rows until the next rowspan row
+    rows_below_until_next_rowspan = []
+
+    # Get the sorted row indices that have rowspan elements
+    sorted_row_indices = sorted(rowspan_dict.keys())
+
+    # Calculate rows below each rowspan row until the next rowspan row
+    for i in range(len(sorted_row_indices)):
+        current_row = sorted_row_indices[i]
+        if i < len(sorted_row_indices) - 1:
+            next_row = sorted_row_indices[i + 1]
+            rows_below = next_row - current_row - 1
+        else:
+            rows_below = len(rows) - current_row - 1
+        rows_below_until_next_rowspan.append((current_row, rows_below))
+
+    # Detect rows where rowspan value is incorrect
+    rows_with_bad_rowspan = []
+    for row_index, rows_below in rows_below_until_next_rowspan:
+        if row_index in rowspan_dict:
+            for cell_index, rowspan_value in rowspan_dict[row_index]:
+                if rowspan_value - 1 < rows_below:
+                    print(f"Row {row_index} has a large rowspan value: {rowspan_value}")
+                    rows_with_bad_rowspan.append(row_index)
+                    break
+
+    # Modify the HTML table to adjust the rowspan attributes
+    for row_index in rows_with_bad_rowspan:
+        if row_index in rowspan_dict:
+            for cell_index, rowspan_value in rowspan_dict[row_index]:
+                # Find the cell with the rowspan attribute
+                cell = rows[row_index].find_all(['td', 'th'])[cell_index]
+                # Remove the rowspan attribute
+                del cell['rowspan']
+                # Find the next row and assign the rowspan value
+                next_row_index = row_index + 1
+                if next_row_index < len(rows):
+                    next_row_cells = rows[next_row_index].find_all(['td', 'th'])
+                    if len(next_row_cells) > cell_index:
+                        next_row_cell = next_row_cells[cell_index]
+                        next_row_cell['rowspan'] = rowspan_value
+                    else:
+                        # Create a new cell if it does not exist
+                        new_cell = soup.new_tag(cell.name)
+                        new_cell['rowspan'] = rowspan_value
+                        new_cell.string = cell.string
+                        rows[next_row_index].append(new_cell)
+
+    # Return the modified HTML table
+    return str(soup)
+
+
+def merge_rows_with_rowspan(html):
+    # Parse the HTML table using BeautifulSoup
+    soup = BeautifulSoup(html, 'html.parser')
+
+    # Extract the header
+    thead = soup.find('thead')
+
+    # Find all rows
+    rows = soup.find_all('tr')
+
+    result = []
+    i = 0
+
+    while i < len(rows):
+        row = rows[i]
+        # Check if any td in the row has a rowspan attribute
+        for td in row.find_all('td'):
+            if td.has_attr('rowspan'):
+                rowspan_value = int(td['rowspan'])
+                result.append(row)
+
+                skip_concatenation = False
+                concatenation_pairs = []
+
+                # Add rows below the current row based on the rowspan number
+                for j in range(1, rowspan_value):
+                    if i + j < len(rows):
+                        below_row = rows[i + j]
+
+                        # Compare cells
+                        row_cells = row.find_all('td')
+                        below_row_cells = below_row.find_all('td')
+                        min_length = min(len(row_cells), len(below_row_cells))
+
+                        for k in range(min_length):
+                            if is_numeric(row_cells[k].get_text(strip=True)) and is_numeric(below_row_cells[k].get_text(strip=True)):
+                                skip_concatenation = True
+                                break
+                            else:
+                                concatenation_pairs.append((row_cells[k], below_row_cells[k]))
+
+                        if skip_concatenation:
+                            result.append(below_row)
+
+                if not skip_concatenation:
+                    for row_cell, below_row_cell in concatenation_pairs:
+                        concatenated_text = (row_cell.get_text(strip=True) + ' ' + below_row_cell.get_text(strip=True)).strip()
+                        row_cell.string = concatenated_text
+
+                i += rowspan_value - 1  # Skip the rows that have been added
+                break
+            else:
+                result.append(row)
+                break
+        i += 1
+
+    # Convert result list of rows back to an HTML table string
+    new_table_soup = BeautifulSoup(f'<table>{str(thead)}</table>', 'html.parser')
+    tbody = new_table_soup.new_tag('tbody')
+    new_table_soup.table.append(tbody)
+    for row in result:
+        for td in row.find_all('td'):
+            if td.has_attr('rowspan'):
+                del td['rowspan']
+        tbody.append(row)
+
+    return str(new_table_soup.table)
+
+
+def detect_and_remove_junk_columns(html_table, target_columns, debug=False):
+    # Wrap the HTML string in a StringIO object
+    html_buffer = StringIO(html_table)
+
+    # Read the HTML table
+    df = pd.read_html(html_buffer)[0]
+
+    model = SentenceTransformer('all-mpnet-base-v2')
+
+    # Get the column names of the dataframe
+    column_names = df.columns.tolist()
+
+    # Calculate the similarity of each column name to the target column names
+    target_embeddings = model.encode(target_columns)
+    column_embeddings = model.encode(column_names)
+
+    # Initialize a dictionary to store the similarity scores
+    similarity_scores = {}
+
+    # Identify junk columns based on similarity threshold
+    junk_columns = []
+    similarity_threshold = 0.5  # Adjust this threshold as needed
+
+    for idx, col_embedding in enumerate(column_embeddings):
+        similarities = util.pytorch_cos_sim(col_embedding, target_embeddings)[0]
+        max_similarity = max(similarities)
+        max_similarity_idx = similarities.argmax().item()  # Get the index of the max similarity
+        similarity_scores[column_names[idx]] = (
+        max_similarity.item(), target_columns[max_similarity_idx])  # Store similarity score and target column name
+        if max_similarity < similarity_threshold:
+            junk_columns.append(column_names[idx])
+
+    if debug:
+        # Print the similarity scores for debugging purposes
+        for column, (score, target_col) in similarity_scores.items():
+            print(f"Column: {column}, Similarity: {score:.4f}, Target Column: {target_col}")
+
+    # Handle junk columns by concatenating their values to the nearest column on the left
+    for junk_col in junk_columns:
+        junk_col_index = column_names.index(junk_col)
+        if junk_col_index > 0:
+            nearest_col = column_names[junk_col_index - 1]
+            df[nearest_col] = df.apply(
+                lambda row: str(row[junk_col]) if pd.isna(row[nearest_col]) and pd.notna(row[junk_col])
+                else (str(row[nearest_col]) + ' ' + str(row[junk_col])) if pd.notna(row[junk_col])
+                else row[nearest_col],
+                axis=1
+            )
+        df.drop(columns=[junk_col], inplace=True)
+
+    # Replace any remaining NaN values with empty strings
+    df = df.fillna('')
+
+    if debug:
+        print(f"Junk columns: {junk_columns}")
+        print(df.to_string())
+
+    # Convert the result into an HTML table
+    html_table = df.to_html(index=False)
+
+    if debug:
+        print(html_table)
+
+    return html_table
+
+
+def clean_html_table_header_names(html_table: str) -> str:
+    """
+        Cleans the headers of an HTML table by removing junk characters and returns the updated HTML as a string.
+
+        Parameters:
+        html (str): The HTML content containing the table.
+
+        Returns:
+        str: The updated HTML table with cleaned headers.
+        """
+    # Parse the HTML table
+    soup = BeautifulSoup(html_table, "html.parser")
+    table = soup.find("table")
+
+    # Extract the headers and clean them
+    headers = table.find_all("th")
+    for th in headers:
+        clean_header = re.sub(r"[^a-zA-Z0-9\s]", "", th.get_text())
+        th.string.replace_with(clean_header)
+
+    html_table = str(soup)
+
+    return html_table
+
+
+def is_numeric(value):
+    # Check if the value is numeric
+    return bool(re.match(r'^\d+(?:,\d{3})*(?:\.\d+)?$', value))

sparrow_parse/extractor/html_extractor.py

@@ -0,0 +1,247 @@
+from rich import print
+from sentence_transformers import SentenceTransformer, util
+from bs4 import BeautifulSoup
+import json
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from extractor_helper import merge_html_table_headers
+from extractor_helper import clean_html_table_header_names
+import re
+
+
+class HTMLExtractor(object):
+    def __init__(self):
+        pass
+
+    def read_data(self, target_columns, data, column_keywords=None, group_by_rows=True, update_targets=False,
+                  local=True, debug=False):
+        answer = {}
+
+        json_result, targets_unprocessed = [], []
+
+        i = 0
+        for table in data:
+            if not target_columns:
+                break
+
+            i += 1
+            json_result, targets_unprocessed = self.read_data_from_table(target_columns, table, column_keywords,
+                                                                         group_by_rows, local, debug)
+            answer = self.add_answer_section(answer, "items" + str(i), json_result)
+
+            if update_targets:
+                target_columns = targets_unprocessed
+
+        answer = self.format_json_output(answer)
+
+        return answer, targets_unprocessed
+
+    def read_data_from_table(self, target_columns, data, column_keywords=None, group_by_rows=True, local=True, debug=False):
+        data = self.invoke_pipeline_step(
+            lambda: merge_html_table_headers(data, column_keywords, debug),
+            "Merging HTML table headers...",
+            local
+        )
+
+        data = self.invoke_pipeline_step(
+            lambda: clean_html_table_header_names(data),
+            "Cleaning HTML table headers...",
+            local
+        )
+
+        columns = self.get_table_column_names(data)
+
+        if debug:
+            print("\n")
+            print(f"Columns: {columns}")
+            print(f"Target columns: {target_columns}")
+
+        indices, targets, targets_unprocessed = self.invoke_pipeline_step(
+            lambda: self.calculate_similarity(columns, target_columns, debug),
+            "Calculating cosine similarity between columns and target values...",
+            local
+        )
+
+        if debug:
+            print(f"Unprocessed targets: {targets_unprocessed}")
+
+        # Extracting data
+        extracted_data = self.invoke_pipeline_step(
+            lambda: self.extract_columns_from_table(data, indices, targets, group_by_rows),
+            "Extracting data from the table...",
+            local
+        )
+
+        json_result = self.convert_to_json(extracted_data)
+
+        return json_result, targets_unprocessed
+
+    def calculate_similarity(self, columns, target_columns, debug):
+        model = SentenceTransformer('all-mpnet-base-v2')
+
+        # Compute embeddings for columns and target values
+        column_embeddings = model.encode(columns)
+        target_embeddings = model.encode(target_columns)
+
+        # List to store indices of the most similar columns
+        most_similar_indices = {}
+        targets_unprocessed = []
+
+        # Calculate cosine similarity between each column and target value
+        similarity_scores = util.pytorch_cos_sim(column_embeddings, target_embeddings)
+
+        # Find the most similar column for each target value and provide the order ID
+        for idx, target in enumerate(target_columns):
+            similarities = similarity_scores[:, idx]
+            most_similar_idx = similarities.argmax().item()
+            most_similar_column = columns[most_similar_idx]
+            similarity_score = similarities[most_similar_idx].item()
+            if similarity_score > 0.3:
+                if most_similar_idx in most_similar_indices:
+                    if similarity_score > most_similar_indices[most_similar_idx][1]:
+                        targets_unprocessed.append(most_similar_indices[most_similar_idx][0])
+                        most_similar_indices[most_similar_idx] = (target, similarity_score)
+                    else:
+                        targets_unprocessed.append(target)
+                else:
+                    most_similar_indices[most_similar_idx] = (target, similarity_score)
+            else:
+                targets_unprocessed.append(target)
+            if debug:
+                print(
+                    f"The most similar column to '{target}' is '{most_similar_column}' with a similarity score of {similarity_score:.4f} and order ID {most_similar_idx}")
+
+        most_similar_indices = dict(sorted(most_similar_indices.items()))
+
+        indices = []
+        targets = []
+
+        for idx, (target, _) in most_similar_indices.items():
+            indices.append(idx)
+            targets.append(target)
+
+        if debug:
+            print()
+            for idx, (target, score) in most_similar_indices.items():
+                print(f"Target: '{target}', Column: '{columns[idx]}', Column ID: {idx}, Score: {score:.4f}")
+            print()
+
+        return indices, targets, targets_unprocessed
+
+    def extract_columns_from_table(self, html_table, column_ids, target_columns, group_by_rows=False):
+        soup = BeautifulSoup(html_table, 'html.parser')
+        table = soup.find('table')
+
+        if group_by_rows:
+            # Initialize a list to store each row's data as a dictionary
+            extracted_data = []
+        else:
+            # Initialize the extracted data with custom column names
+            extracted_data = {target_columns[i]: [] for i in range(len(column_ids))}
+
+        # Extract row information
+        rows = table.find_all('tr')
+
+        for row in rows:
+            # Skip the header row
+            if row.find_all('th'):
+                continue
+
+            cells = row.find_all('td')
+            if cells:  # Ensure the row contains data cells
+                if group_by_rows:
+                    row_data = {}
+                    for idx, col_id in enumerate(column_ids):
+                        value = cells[col_id].text.strip() if col_id < len(cells) else ''
+                        value = value.replace('|', '').strip()
+                        row_data[target_columns[idx]] = value
+                    extracted_data.append(row_data)
+                else:
+                    for idx, col_id in enumerate(column_ids):
+                        value = cells[col_id].text.strip() if col_id < len(cells) else ''
+                        value = value.replace('|', '').strip()
+                        extracted_data[target_columns[idx]].append(value)
+
+        return extracted_data
+
+    def convert_to_json(self, extracted_data):
+        return json.dumps(extracted_data, indent=4)
+
+    def get_table_column_names(self, html_table):
+        """
+        Extract column names from an HTML table.
+
+        Args:
+        html_table (str): The HTML content of the table.
+
+        Returns:
+        list: A list of column names.
+        """
+        # Parse the HTML content using BeautifulSoup with html.parser
+        soup = BeautifulSoup(html_table, 'html.parser')
+
+        # Find the <thead> tag
+        thead = soup.find('thead')
+
+        # Extract column names into a list
+        column_names = [th.get_text() for th in thead.find_all('th')]
+
+        return column_names
+
+    def invoke_pipeline_step(self, task_call, task_description, local):
+        if local:
+            with Progress(
+                    SpinnerColumn(),
+                    TextColumn("[progress.description]{task.description}"),
+                    transient=False,
+            ) as progress:
+                progress.add_task(description=task_description, total=None)
+                ret = task_call()
+        else:
+            print(task_description)
+            ret = task_call()
+
+        return ret
+
+    def add_answer_section(self, answer, section_name, answer_table):
+        if not isinstance(answer, dict):
+            raise ValueError("The answer should be a dictionary.")
+
+        # Parse answer_table if it is a JSON string
+        if isinstance(answer_table, str):
+            answer_table = json.loads(answer_table)
+
+        answer[section_name] = answer_table
+        return answer
+
+    def format_json_output(self, answer):
+        formatted_json = json.dumps(answer, indent=4)
+        formatted_json = formatted_json.replace('", "', '",\n"')
+        formatted_json = formatted_json.replace('}, {', '},\n{')
+        return formatted_json
+
+
+if __name__ == "__main__":
+    # with open('../data/invoice_1_table.txt', 'r') as file:
+    #     file_content = file.read()
+    #
+    # file_content = file_content.strip()[1:-1].strip()
+    # data_list = re.split(r"',\s*'", file_content)
+    # data_list = [item.strip(" '") for item in data_list]
+
+    extractor = HTMLExtractor()
+
+    # answer, targets_unprocessed = extractor.read_data(
+    #     ['description', 'qty', 'net_price', 'net_worth', 'vat', 'gross_worth'],
+    #     # ['transaction_date', 'value_date', 'description', 'cheque', 'withdrawal', 'deposit', 'balance',
+    #     #  'deposits', 'account_number', 'od_limit', 'currency_balance', 'sgd_balance', 'maturity_date'],
+    #     data_list,
+    #     None,
+    #     # ['deposits', 'account_number', 'od_limit', 'currency_balance', 'sgd_balance', 'transaction_date',
+    #     #  'value_date', 'description', 'cheque', 'withdrawal', 'deposit', 'balance', 'maturity_date'],
+    #     True,
+    #     True,
+    #     True,
+    #     True)
+    #
+    # print(answer)
+    # print(targets_unprocessed)

sparrow_parse/extractor/markdown_processor.py

@@ -0,0 +1,137 @@
+import pymupdf4llm
+import pandas as pd
+import re
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich import print
+from bs4 import BeautifulSoup
+
+
+class MarkdownProcessor(object):
+    def __init__(self):
+        pass
+
+    def extract_data(self, file_path, options, local=True, debug=False):
+        markdown_text = self.invoke_pipeline_step(
+            lambda: pymupdf4llm.to_markdown(file_path),
+            "Extracting markdown text from the document...",
+            local
+        )
+
+        content, table_content = self.invoke_pipeline_step(
+            lambda: self.load_text_data(markdown_text, options),
+            "Loading text data...",
+            local
+        )
+
+        if debug:
+            print("Data extracted from the document:")
+            print(content)
+            print("\n")
+            print("Table content extracted from the document:")
+            if table_content:
+                print(len(table_content))
+            print(table_content)
+
+        return content, table_content
+
+    def load_text_data(self, markdown_text, options):
+        content, table_content = None, None
+
+        if options is None:
+            content = markdown_text
+
+        if options and "tables" in options and "markdown" in options:
+            content = self.extract_form_data(markdown_text)
+            table_content = self.extract_tables(markdown_text)
+
+        return content, table_content
+
+    def extract_form_data(self, markdown_text):
+        return markdown_text
+
+    def extract_tables(self, markdown_text):
+        # Regular expression to match markdown tables
+        table_pattern = re.compile(r'(\|.+\|\n\|[-| ]+\|\n(?:\|.*\|\n)*?)(?=\|.*TOTAL)', re.MULTILINE)
+
+        # Find all tables in the markdown text
+        tables = table_pattern.findall(markdown_text)
+
+        html_tables = []
+        for table_text in tables:
+            # Split the table into lines
+            lines = table_text.strip().split('\n')
+
+            # Extract headers and rows
+            headers = [self.clean_column_name(header.strip()) for header in lines[0].split('|') if header]
+            rows = []
+            for line in lines[2:]:  # Skip header and separator lines
+                row = [cell.strip() for cell in line.split('|') if cell]
+                rows.append(row)
+
+            # Convert to Pandas DataFrame
+            df = pd.DataFrame(rows, columns=headers)
+
+            # Convert DataFrame to HTML and append to the list
+            html_table = df.to_html(index=False)
+            if self.table_has_header(html_table):
+                html_tables.append(html_table)
+
+            return html_tables
+
+    def clean_column_name(self, name):
+        """
+        Cleans the column name by removing spaces if the name is a single word with spaces between letters.
+
+        Args:
+        name (str): The column name to clean.
+
+        Returns:
+        str: The cleaned column name.
+        """
+        # Check if the name contains only letters and spaces
+        if all(char.isalpha() or char.isspace() for char in name):
+            # Check if it is a single word with spaces between letters
+            parts = name.split()
+            if len(parts) > 1 and all(len(part) == 1 for part in parts):
+                return ''.join(parts)
+        return name
+
+    def invoke_pipeline_step(self, task_call, task_description, local):
+        if local:
+            with Progress(
+                    SpinnerColumn(),
+                    TextColumn("[progress.description]{task.description}"),
+                    transient=False,
+            ) as progress:
+                progress.add_task(description=task_description, total=None)
+                ret = task_call()
+        else:
+            print(task_description)
+            ret = task_call()
+
+        return ret
+
+    def table_has_header(self, table_html):
+        soup = BeautifulSoup(table_html, 'html.parser')
+        table = soup.find('table')
+
+        # Check if the table contains a <thead> tag
+        if table.find('thead'):
+            return True
+
+        # Check if the table contains any <th> tags inside the table (in case there's no <thead>)
+        if table.find_all('th'):
+            return True
+
+        return False
+
+
+if __name__ == "__main__":
+    processor = MarkdownProcessor()
+
+    # content, table_content = processor.extract_data(
+    #     '/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/invoice_1.pdf',
+    #     ['tables', 'markdown'],
+    #     True,
+    #     True)
+

sparrow_parse/extractor/unstructured_processor.py

@@ -0,0 +1,179 @@
+import tempfile
+import os
+from unstructured.partition.pdf import partition_pdf
+from unstructured.partition.image import partition_image
+import json
+from unstructured.staging.base import elements_to_json
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich import print
+from bs4 import BeautifulSoup
+
+
+class UnstructuredProcessor(object):
+    def __init__(self):
+        pass
+
+    def extract_data(self, file_path, strategy, model_name, options, local=True, debug=False):
+        # Extracts the elements from the PDF
+        elements = self.invoke_pipeline_step(
+            lambda: self.process_file(file_path, strategy, model_name),
+            "Extracting elements from the document...",
+            local
+        )
+
+        if debug:
+            new_extension = 'json'  # You can change this to any extension you want
+            new_file_path = self.change_file_extension(file_path, new_extension)
+
+            content, table_content = self.invoke_pipeline_step(
+                lambda: self.load_text_data(elements, new_file_path, options),
+                "Loading text data...",
+                local
+            )
+        else:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                temp_file_path = os.path.join(temp_dir, "file_data.json")
+
+                content, table_content = self.invoke_pipeline_step(
+                    lambda: self.load_text_data(elements, temp_file_path, options),
+                    "Loading text data...",
+                    local
+                )
+
+        if debug:
+            print("Data extracted from the document:")
+            print(content)
+            print("\n")
+            print("Table content extracted from the document:")
+            if table_content:
+                print(len(table_content))
+            print(table_content)
+
+        return content, table_content
+
+    def process_file(self, file_path, strategy, model_name):
+        elements = None
+
+        if file_path.lower().endswith('.pdf'):
+            elements = partition_pdf(
+                filename=file_path,
+                strategy=strategy,
+                infer_table_structure=True,
+                hi_res_model_name=model_name,
+                languages=['en']
+            )
+        elif file_path.lower().endswith(('.jpg', '.jpeg', '.png')):
+            elements = partition_image(
+                filename=file_path,
+                strategy=strategy,
+                infer_table_structure=True,
+                hi_res_model_name=model_name,
+                languages=['en']
+            )
+
+        return elements
+
+    def change_file_extension(self, file_path, new_extension, suffix=None):
+        # Check if the new extension starts with a dot and add one if not
+        if not new_extension.startswith('.'):
+            new_extension = '.' + new_extension
+
+        # Split the file path into two parts: the base (everything before the last dot) and the extension
+        # If there's no dot in the filename, it'll just return the original filename without an extension
+        base = file_path.rsplit('.', 1)[0]
+
+        # Concatenate the base with the new extension
+        if suffix is None:
+            new_file_path = base + new_extension
+        else:
+            new_file_path = base + "_" + suffix + new_extension
+
+        return new_file_path
+
+    def load_text_data(self, elements, file_path, options):
+        elements_to_json(elements, filename=file_path)
+
+        content, table_content = None, None
+
+        if options is None:
+            content = self.process_json_file(file_path)
+
+        if options and "tables" in options and "html" in options:
+            content = self.process_json_file(file_path, "form")
+
+            table_content = self.process_json_file(file_path, "table")
+
+        return content, table_content
+
+    def process_json_file(self, file_path, option=None):
+        # Read the JSON file
+        with open(file_path, 'r') as file:
+            data = json.load(file)
+
+        # Iterate over the JSON data and extract required elements
+        extracted_elements = []
+        for entry in data:
+            if entry["type"] == "Table" and (option is None or option == "table" or option == "form"):
+                table_data = entry["metadata"]["text_as_html"]
+                if option == "table" and self.table_has_header(table_data):
+                    extracted_elements.append(table_data)
+                if option is None or option == "form":
+                    extracted_elements.append(table_data)
+            elif entry["type"] == "Title" and (option is None or option == "form"):
+                extracted_elements.append(entry["text"])
+            elif entry["type"] == "NarrativeText" and (option is None or option == "form"):
+                extracted_elements.append(entry["text"])
+            elif entry["type"] == "UncategorizedText" and (option is None or option == "form"):
+                extracted_elements.append(entry["text"])
+            elif entry["type"] == "ListItem" and (option is None or option == "form"):
+                extracted_elements.append(entry["text"])
+            elif entry["type"] == "Image" and (option is None or option == "form"):
+                extracted_elements.append(entry["text"])
+
+        if option is None or option == "form":
+            # Convert list to single string with two new lines between each element
+            extracted_data = "\n\n".join(extracted_elements)
+            return extracted_data
+
+        return extracted_elements
+
+    def invoke_pipeline_step(self, task_call, task_description, local):
+        if local:
+            with Progress(
+                    SpinnerColumn(),
+                    TextColumn("[progress.description]{task.description}"),
+                    transient=False,
+            ) as progress:
+                progress.add_task(description=task_description, total=None)
+                ret = task_call()
+        else:
+            print(task_description)
+            ret = task_call()
+
+        return ret
+
+    def table_has_header(self, table_html):
+        soup = BeautifulSoup(table_html, 'html.parser')
+        table = soup.find('table')
+
+        # Check if the table contains a <thead> tag
+        if table.find('thead'):
+            return True
+
+        # Check if the table contains any <th> tags inside the table (in case there's no <thead>)
+        if table.find_all('th'):
+            return True
+
+        return False
+
+
+if __name__ == "__main__":
+    processor = UnstructuredProcessor()
+    # content, table_content = processor.extract_data(
+    #     '/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/invoice_1.pdf',
+    #     'hi_res',
+    #     'yolox',
+    #     ['tables', 'html'],
+    #     True,
+    #     True)
+

sparrow_parse/temp.py

@@ -0,0 +1,28 @@
+# content, table_content = processor.extract_data(
+    #     '/Users/andrejb/Documents/work/schreiber/invoice_data/test/2618407.pdf',
+    #     'hi_res',
+    #     'yolox',
+    #     # 'detectron2_onnx',
+    #     ['tables', 'html'],
+    #     True,
+    #     True)
+
+    # content, table_content = processor.extract_data(
+    #     '/Users/andrejb/Documents/work/epik/bankstatement/OCBC_1_1.pdf',
+    #     'hi_res',
+    #     'yolox',
+    #     ['tables', 'html'],
+    #     True,
+    #     True)
+
+
+# content, table_content = processor.extract_data(
+#     '/Users/andrejb/Documents/work/schreiber/invoice_data/test/2618407.pdf',
+#     ['tables', 'markdown'],
+#     True,
+#     True)
+# content, table_content = processor.extract_data(
+#     '/Users/andrejb/Documents/work/epik/bankstatement/POSB_2_1.pdf',
+#     ['tables', 'markdown'],
+#     True,
+#     True)

{sparrow_parse-0.1.9.dist-info → sparrow_parse-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sparrow-parse
-Version: 0.1.9
+Version: 0.2.0
 Summary: Sparrow Parse is a Python package for parsing and extracting information from documents.
 Home-page: https://github.com/katanaml/sparrow/tree/main/sparrow-data/parse
 License: GPL-3.0
@@ -15,10 +15,13 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: Software Development
+Requires-Dist: pymupdf4llm (==0.0.5)
 Requires-Dist: rich (>=13.7.1,<14.0.0)
+Requires-Dist: sentence-transformers (==3.0.1)
 Requires-Dist: torch (==2.2.2)
-Requires-Dist: unstructured-inference (==0.7.29)
-Requires-Dist: unstructured[all-docs] (==0.13.6)
+Requires-Dist: transformers (==4.41.2)
+Requires-Dist: unstructured-inference (==0.7.33)
+Requires-Dist: unstructured[all-docs] (==0.14.5)
 Project-URL: Repository, https://github.com/katanaml/sparrow
 Description-Content-Type: text/markdown
 
@@ -26,7 +29,7 @@ Description-Content-Type: text/markdown
 
 ## Description
 
-This module implements Sparrow Parse [library](https://pypi.org/project/sparrow-parse/) with helpful methods for data pre-processing.
+This module implements Sparrow Parse [library](https://pypi.org/project/sparrow-parse/) with helpful methods for data pre-processing, parsing and extracting information.
 
 ## Install
 
@@ -34,22 +37,97 @@ This module implements Sparrow Parse [library](https://pypi.org/project/sparrow-
 pip install sparrow-parse
 ```
 
-## Use
+## Pre-processing
 
-Import
+### Unstructured
 
 ```
-from sparrow_parse.pdf.pdf_processor import PDFProcessor
+from sparrow_parse.extractor.unstructured_processor import UnstructuredProcessor
+
+processor = UnstructuredProcessor()
+
+content, table_content = processor.extract_data(
+        file_path,  # file to process
+        strategy,  # data processing strategy supported by unstructured
+        model_name,  # model supported by unstructured
+        options,  # table extraction into HTML format
+        local,  # True if running from CLI, or False if running from FastAPI
+        debug)  # Debug
 ```
 
-Usage
+Example:
+
+*file_path* - `/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/invoice_1.pdf`
+
+*strategy* - `hi_res`
+
+*model_name* - `yolox`
+
+*options* - `['tables', 'html']`
+
+*local* - `True`
+
+*debug* - `True`
+
+### Markdown
 
 ```
-processor = PDFProcessor()
-result = processor.process_file(file_path, strategy, model_name)
+from sparrow_parse.extractor.markdown_processor import MarkdownProcessor
+
+processor = MarkdownProcessor()
+
+content, table_content = processor.extract_data(
+        file_path,  # file to process
+        options,  # table extraction into HTML format
+        local,  # True if running from CLI, or False if running from FastAPI
+        debug)  # Debug
 ```
 
-Build for development
+Example:
+
+*file_path* - `/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/invoice_1.pdf`
+
+*options* - `['tables', 'markdown']`
+
+*local* - `True`
+
+*debug* - `True`
+
+## Parsing and extraction
+
+```
+from sparrow_parse.extractor.html_extractor import HTMLExtractor
+
+extractor = HTMLExtractor()
+
+answer, targets_unprocessed = extractor.read_data(
+        target_columns,  # list of table columns data to fetch
+        data, # list of HTML tables
+        column_keywords,  # list of valid column names, can be empty. Useful to filter junk content
+        group_by_rows,  # JSON result grouping
+        update_targets,  # Set to true, if page contains multiple tables with the same columns
+        local,  # True if running from CLI, or False if running from FastAPI
+        debug)  # Debug
+
+```
+
+Example:
+
+*target_columns* - `['description', 'qty', 'net_price', 'net_worth', 'vat', 'gross_worth']`
+
+*data* - `list of HTML tables`
+
+*column_keywords* - `None`
+
+*group_by_rows* - `True`
+
+*update_targets* - `True`
+
+*local* - `True`
+
+*debug* - `True`
+
+## Library build
 
 ```
 poetry build

sparrow_parse-0.2.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+sparrow_parse/__init__.py,sha256=clN3TnyYyt5T_wUJLVBmutS2kYLLBET4JFB2QXnRm2Q,21
+sparrow_parse/__main__.py,sha256=Xs1bpJV0n08KWOoQE34FBYn6EBXZA9HIYJKrE4ZdG78,153
+sparrow_parse/data/invoice_1_table.txt,sha256=dsWEASxlVNidpTCQDowCM7SjaUzSqwx7DuydTfaQ7xI,1115
+sparrow_parse/extractor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sparrow_parse/extractor/extractor_helper.py,sha256=fvA7iuGdpW_WFc4jkALBzQbACqlE5x_K9ScW-E6RCoY,13357
+sparrow_parse/extractor/html_extractor.py,sha256=XGmaceTuny185iGiLVikEZ8ksLex6mRvQhw-McB84CU,9381
+sparrow_parse/extractor/markdown_processor.py,sha256=dC2WUdA-v2psh7oytruftxYkXdQi72FoEYxF30ROuO0,4506
+sparrow_parse/extractor/unstructured_processor.py,sha256=7eFIZ6VeEkIslcJyibyix1qFPHGtIwnvNxdRItGz7V8,6755
+sparrow_parse/temp.py,sha256=xiRDPkv_fsM9xCcW29TU0LushgYjMaN9_Cwur6RvY1A,859
+sparrow_parse-0.2.0.dist-info/METADATA,sha256=jH_xt6142QixaARXk5EMp01133YRg8tPvKzAfsjQVGE,5582
+sparrow_parse-0.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+sparrow_parse-0.2.0.dist-info/entry_points.txt,sha256=H507qotwq3VX4lv5pY9MZYtupKNE1RRb8gEQucPiGi0,52
+sparrow_parse-0.2.0.dist-info/RECORD,,

sparrow_parse/extractor/file_processor.py

@@ -1,143 +0,0 @@
-import tempfile
-import os
-from unstructured.partition.pdf import partition_pdf
-from unstructured.partition.image import partition_image
-import json
-from unstructured.staging.base import elements_to_json
-from rich.progress import Progress, SpinnerColumn, TextColumn
-
-
-class FileProcessor(object):
-    def __init__(self):
-        pass
-
-    def extract_data(self, file_path, strategy, model_name, options, local=True, debug=False):
-        # check if string options contains word table
-        extract_tables = False
-        if options is not None and "tables" in options:
-            extract_tables = True
-
-        # Extracts the elements from the PDF
-        elements = self.invoke_pipeline_step(
-            lambda: self.process_file(file_path, strategy, model_name),
-            "Extracting elements from the document...",
-            local
-        )
-
-        if debug:
-            new_extension = 'json'  # You can change this to any extension you want
-            new_file_path = self.change_file_extension(file_path, new_extension)
-
-            content = self.invoke_pipeline_step(
-                lambda: self.load_text_data(elements, new_file_path, extract_tables),
-                "Loading text data...",
-                local
-            )
-        else:
-            with tempfile.TemporaryDirectory() as temp_dir:
-                temp_file_path = os.path.join(temp_dir, "file_data.json")
-
-                content = self.invoke_pipeline_step(
-                    lambda: self.load_text_data(elements, temp_file_path, extract_tables),
-                    "Loading text data...",
-                    local
-                )
-
-        return content
-
-    def process_file(self, file_path, strategy, model_name):
-        elements = None
-
-        if file_path.lower().endswith('.pdf'):
-            elements = partition_pdf(
-                filename=file_path,
-                strategy=strategy,
-                infer_table_structure=True,
-                model_name=model_name
-            )
-        elif file_path.lower().endswith(('.jpg', '.jpeg', '.png')):
-            elements = partition_image(
-                filename=file_path,
-                strategy=strategy,
-                infer_table_structure=True,
-                model_name=model_name
-            )
-
-        return elements
-
-    def change_file_extension(self, file_path, new_extension):
-        # Check if the new extension starts with a dot and add one if not
-        if not new_extension.startswith('.'):
-            new_extension = '.' + new_extension
-
-        # Split the file path into two parts: the base (everything before the last dot) and the extension
-        # If there's no dot in the filename, it'll just return the original filename without an extension
-        base = file_path.rsplit('.', 1)[0]
-
-        # Concatenate the base with the new extension
-        new_file_path = base + new_extension
-
-        return new_file_path
-
-    def load_text_data(self, elements, file_path, extract_tables):
-        elements_to_json(elements, filename=file_path)
-        text_file = self.process_json_file(file_path, extract_tables)
-
-        with open(text_file, 'r') as file:
-            content = file.read()
-
-        return content
-
-    def process_json_file(self, input_data, extract_tables):
-        # Read the JSON file
-        with open(input_data, 'r') as file:
-            data = json.load(file)
-
-        # Iterate over the JSON data and extract required table elements
-        extracted_elements = []
-        for entry in data:
-            if entry["type"] == "Table":
-                extracted_elements.append(entry["metadata"]["text_as_html"])
-            elif entry["type"] == "Title" and extract_tables is False:
-                extracted_elements.append(entry["text"])
-            elif entry["type"] == "NarrativeText" and extract_tables is False:
-                extracted_elements.append(entry["text"])
-            elif entry["type"] == "UncategorizedText" and extract_tables is False:
-                extracted_elements.append(entry["text"])
-
-        # Write the extracted elements to the output file
-        new_extension = 'txt'  # You can change this to any extension you want
-        new_file_path = self.change_file_extension(input_data, new_extension)
-        with open(new_file_path, 'w') as output_file:
-            for element in extracted_elements:
-                output_file.write(element + "\n\n")  # Adding two newlines for separation
-
-        return new_file_path
-
-    def invoke_pipeline_step(self, task_call, task_description, local):
-        if local:
-            with Progress(
-                    SpinnerColumn(),
-                    TextColumn("[progress.description]{task.description}"),
-                    transient=False,
-            ) as progress:
-                progress.add_task(description=task_description, total=None)
-                ret = task_call()
-        else:
-            print(task_description)
-            ret = task_call()
-
-        return ret
-
-
-# if __name__ == "__main__":
-#     processor = FileProcessor()
-#     content = processor.extract_data('/Users/andrejb/infra/shared/katana-git/sparrow/sparrow-ml/llm/data/invoice_1.pdf',
-#                                      'hi_res',
-#                                      'yolox',
-#                                      'tables',
-#                                      False,
-#                                      True)
-#     processor.extract_data("/Users/andrejb/Documents/work/lifung/lemming_test/C16E150001_SUPINV.pdf")
-#     processor.extract_data("/Users/andrejb/Documents/work/epik/bankstatement/OCBC_1_single.pdf")
-#     print(content)

sparrow_parse-0.1.9.dist-info/RECORD

@@ -1,8 +0,0 @@
-sparrow_parse/__init__.py,sha256=m1D6fscvvsMhq5HVNKw7kP5M8AqEzQm1ekrn_nLQF1M,21
-sparrow_parse/__main__.py,sha256=Xs1bpJV0n08KWOoQE34FBYn6EBXZA9HIYJKrE4ZdG78,153
-sparrow_parse/extractor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sparrow_parse/extractor/file_processor.py,sha256=OrDxFJVEhy_4pCBxknehAM7fxgSlWgUJ0jeTEegHRxo,5621
-sparrow_parse-0.1.9.dist-info/METADATA,sha256=wK7uOpPqsC1iwZs_d5Hl1KV4DIbw8NnfPU6MBTyF_kA,3428
-sparrow_parse-0.1.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-sparrow_parse-0.1.9.dist-info/entry_points.txt,sha256=H507qotwq3VX4lv5pY9MZYtupKNE1RRb8gEQucPiGi0,52
-sparrow_parse-0.1.9.dist-info/RECORD,,