sparrow-parse 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

sparrow_parse/__init__.py

@@ -1 +1 @@
-__version__ = '0.1.8'
+__version__ = '0.1.10'

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,145 @@
+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)
+    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/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,16 @@
+# 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)

{sparrow_parse-0.1.8.dist-info → sparrow_parse-0.1.10.dist-info}/METADATA

@@ -1,21 +1,27 @@
 Metadata-Version: 2.1
 Name: sparrow-parse
-Version: 0.1.8
+Version: 0.1.10
 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
 Keywords: llm,rag,vision
 Author: Andrej Baranovskij
 Author-email: andrejus.baranovskis@gmail.com
-Requires-Python: >=3.10,<4.0
+Requires-Python: >=3.9,<3.12
 Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Software Development
-Requires-Dist: requests (>=2.31.0,<3.0.0)
+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: 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
 
@@ -36,20 +42,26 @@ pip install sparrow-parse
 Import
 
 ```
-from sparrow_parse.pdf.pdf_processor import PDFProcessor
+from sparrow_parse.extractor.file_processor import FileProcessor
 ```
 
 Usage
 
 ```
-processor = PDFProcessor()
-result = processor.process_file(file_path, strategy, model_name)
+processor = FileProcessor()
+content = processor.extract_data(file_path, strategy, model_name, options, local, debug)
 ```
 
-Build for development
+## Library build
 
 ```
+poetry build
+```
+
+Publish to PyPi
 
+```
+poetry publish
 ```
 
 ## Commercial usage

sparrow_parse-0.1.10.dist-info/RECORD

@@ -0,0 +1,13 @@
+sparrow_parse/__init__.py,sha256=GG5_e12GvnOxWm5sUjputYu0F3PdW-4ThVZhf-KKQHY,22
+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=YnrDYvmBygF-TC3_BlzEMWI6w_PWxvQWEw2zG37AjGw,4894
+sparrow_parse/extractor/unstructured_processor.py,sha256=7eFIZ6VeEkIslcJyibyix1qFPHGtIwnvNxdRItGz7V8,6755
+sparrow_parse/temp.py,sha256=zY77otqPZYvHqoyo12jxKIzr2eTMWI1h1wZnv4N0r3s,497
+sparrow_parse-0.1.10.dist-info/METADATA,sha256=xqQeYGc6OKhBuFG8IgiV9k704ZePAw4tUhuRaifZTNc,3580
+sparrow_parse-0.1.10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+sparrow_parse-0.1.10.dist-info/entry_points.txt,sha256=H507qotwq3VX4lv5pY9MZYtupKNE1RRb8gEQucPiGi0,52
+sparrow_parse-0.1.10.dist-info/RECORD,,

sparrow_parse/pdf/pdf_processor.py

@@ -1,7 +0,0 @@
-class PDFProcessor(object):
-    def __init__(self):
-        pass
-
-    def process_file(self, content):
-        print("Processing file...")
-        return "OK"

sparrow_parse-0.1.8.dist-info/RECORD

@@ -1,8 +0,0 @@
-sparrow_parse/__init__.py,sha256=zemvJ5zjFE6SQT2xmkxc-ZYwNkUTCEX7mz3Epb2qztE,21
-sparrow_parse/__main__.py,sha256=Xs1bpJV0n08KWOoQE34FBYn6EBXZA9HIYJKrE4ZdG78,153
-sparrow_parse/pdf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sparrow_parse/pdf/pdf_processor.py,sha256=hyvOQX_IydRA3z7gQs_g-Ut1hvVHRRxj1_2i-G09-ow,159
-sparrow_parse-0.1.8.dist-info/METADATA,sha256=QTAeFIi-KwyBvSbBrB8wS5WCld3gQ3XfAll4wS4x7Yc,3250
-sparrow_parse-0.1.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-sparrow_parse-0.1.8.dist-info/entry_points.txt,sha256=H507qotwq3VX4lv5pY9MZYtupKNE1RRb8gEQucPiGi0,52
-sparrow_parse-0.1.8.dist-info/RECORD,,