Upload post_process_portfolio_company_json.py

post_process_portfolio_company_json.py

@@ -0,0 +1,387 @@
+import json
+import os
+from fuzzywuzzy import fuzz
+from typing import List, Dict, Any
+import yaml
+import warnings
+import pandas as pd
+
+# Constants
+# PORTFOLIO_COMPANY_LIST_IDENTIFIER = ["portfolio company or platforms", "portfolio company"]
+PORTFOLIO_COMPANY_LIST_IDENTIFIER = ["portfolio company or platforms","\u20acm","$m","Unrealised fair market valuation","Realised proceeds in the period","Portfolio Company or Platforms","portfolio company", "active investment", "realized/unrealized company","Realized Company","Unrealized Company", "quoted/unquoted company", "portfolio investment", "portfolio company"]
+FUZZY_MATCH_THRESHOLD = 70
+EXCLUDE_COMPANY_NAMES = ["total", "subtotal","Total","Investments","Fund"]
+
+
+def get_file_name_without_extension(file_path: str) -> str:
+    """Extract file name without extension from path."""
+    return os.path.splitext(os.path.basename(file_path))[0]
+
+def fuzzy_match(text: str, patterns: List[str], threshold: int = FUZZY_MATCH_THRESHOLD) -> bool:
+    """Check if text fuzzy matches any of the patterns."""
+    text = str(text).lower()
+    for pattern in patterns:
+        if fuzz.partial_ratio(text, pattern.lower()) >= threshold:
+            return True
+    return False
+
+def extract_portfolio_companies_from_table(table_data: Dict) -> List[str]:
+    """Extract company names from a portfolio company table."""
+    companies = []
+    if not table_data.get("table_info"):
+        return companies
+    
+    # Find the company column
+    company_column = None
+    for i, header in enumerate(table_data.get("table_column_header", [])):
+        if fuzzy_match(header, PORTFOLIO_COMPANY_LIST_IDENTIFIER):
+            company_column = i
+            break
+    
+    if company_column is None:
+        return companies
+    
+    # Get the column name that contains companies
+    company_column_name = table_data["table_column_header"][company_column]
+    print("company_column::",company_column)
+    print("cpmpany_column_name::",company_column_name)
+    
+    # Extract companies
+    for row in table_data["table_info"]:
+        if not isinstance(row, dict):
+            continue
+        company_name = str(row.get(company_column_name, "")).strip()
+        if company_name and not fuzzy_match(company_name, EXCLUDE_COMPANY_NAMES):
+            companies.append(company_name)
+    
+    return companies
+
+def get_portfolio_company_list(intermediate_data: List[Dict]) -> List[str]:
+    """Extract portfolio companies from all tables in the document."""
+    portfolio_companies = set()
+    
+    for entry in intermediate_data:
+        if "table_content" not in entry:
+            continue
+        for table in entry["table_content"]:
+            companies = extract_portfolio_companies_from_table(table)
+            portfolio_companies.update(companies)
+    
+    return list(portfolio_companies)
+
+def merge_content_under_same_header(
+    intermediate_data: List[Dict],
+    portfolio_company_list: List[str],
+    start_index: int
+) -> Dict:
+    """
+    Merge content under the same header until next company match is found.
+    Returns merged content and the next index to process.
+    """
+    merged_entry = {
+        "header": intermediate_data[start_index]["header"],
+        "content": intermediate_data[start_index].get("content", ""),
+        "table_content": intermediate_data[start_index].get("table_content", []),
+        "label_name": intermediate_data[start_index]["label_name"],
+        "page_number": intermediate_data[start_index]["page_number"],
+        "pdf_name": intermediate_data[start_index]["pdf_name"]
+    }
+    
+    current_index = start_index + 1
+    while current_index < len(intermediate_data):
+        current_entry = intermediate_data[current_index]
+        
+        # Check if we're still under the same header
+        if current_entry["header"] != merged_entry["header"]:
+            break
+            
+        # Check if current entry matches any portfolio company
+        content_match = any(company in current_entry.get("content", "") 
+                          for company in portfolio_company_list)
+        table_match = False
+        for table in current_entry.get("table_content", []):
+            if extract_portfolio_companies_from_table(table):
+                table_match = True
+                break
+        
+        if content_match or table_match:
+            break
+        
+        # Merge content
+        if "content" in current_entry:
+            if merged_entry["content"]:
+                merged_entry["content"] += "\n" + current_entry["content"]
+            else:
+                merged_entry["content"] = current_entry["content"]
+        
+        # Merge tables
+        if "table_content" in current_entry:
+            merged_entry["table_content"].extend(current_entry["table_content"])
+        
+        current_index += 1
+    
+    return merged_entry, current_index
+
+
+
+def process_table_page_ids(merged_output):
+    """
+    Process the data to update the page_number key by combining its existing values with unique page numbers
+    from table_content metadata, for pages that contain table_content.
+    
+    Args:
+        data (dict): Input data dictionary with page numbers as keys and page content as values.
+    
+    Returns:
+        dict: Modified data with updated page_number key including existing and metadata page numbers.
+    """
+    # Iterate through each page in the data
+    for current_merged_entry in merged_output:
+        # Only process pages that have table_content
+        if 'table_content' in current_merged_entry:
+            # Initialize a set with existing page numbers from the page_number key
+            existing_page_numbers = set(current_merged_entry.get('page_number', '').split(',')) if current_merged_entry.get('page_number') else set()
+            
+            # Add unique page numbers from table_content metadata
+            for table in current_merged_entry['table_content']:
+                if 'metadata' in table and 'table_page_id' in table['metadata']:
+                    existing_page_numbers.add(str(table['metadata']['table_page_id']))
+            
+            # Update the page_number key with sorted, unique page numbers
+            if existing_page_numbers:
+                current_merged_entry['page_number'] = ','.join(sorted(existing_page_numbers, key=int))
+    
+    return merged_output
+
+
+################################################################################################################
+## Below function for more than one occurence of underlying_assets
+
+def merge_portfolio_company_sections(intermediate_data: List[Dict]) -> tuple[List[Dict], List[str], List[str]]:
+    """Merge all content and tables under the same portfolio company header until next company is found.
+    Returns:
+        - merged_output: List of merged document sections
+        - fuzzy_matched_companies: List of companies that were fuzzy matched in headers
+        - portfolio_companies: List of all portfolio companies found in tables
+    """
+    portfolio_companies = get_portfolio_company_list(intermediate_data)
+    print(f"Extracted portfolio companies: {portfolio_companies}")
+    
+    merged_output = []
+    # fuzzy_matched_companies = set()
+    current_chunk = None
+    active_company = None
+    
+    for entry in intermediate_data:
+        # Find all companies in this entry's header
+        # header_companies = []
+        # for company in portfolio_companies:
+        #     if fuzzy_match(entry["header"], [company], threshold=90):
+        #         header_companies.append(company)
+        #         fuzzy_matched_companies.add(company)
+        entry_copy = entry.copy()
+
+        header_companies, fuzzy_matched_companies = match_company_names(entry["header"], portfolio_companies)
+        # print("header_companies::",header_companies)
+        # print("fuzzy_matched_companies::",fuzzy_matched_companies)
+        
+        if header_companies:
+            print("&"*100)
+            print("*"*100)
+            print("entry_header::", entry["header"])
+            print("page number of header::", entry["page_number"])
+
+            print("*"*100)
+            print("header_companies::", header_companies)
+            print("*"*100)
+            
+            # If we have an active chunk, finalize it before starting new one
+            if current_chunk:
+                merged_output.append(current_chunk)
+                current_chunk = None
+                active_company = None
+            
+            # Start new chunk with the first matched company
+            # (in case multiple companies matched, we take the first one)
+            active_company = header_companies[0]
+            current_chunk = {
+                "page_number": entry["page_number"],
+                "pdf_name": entry["pdf_name"],
+                "header": entry["header"],
+                "label_name": entry["label_name"],
+                "content": entry.get("content", ""),
+                "table_content": entry.get("table_content", []),
+                "matched_company": active_company
+            }
+            
+            # If multiple companies matched, create separate chunks for others
+            for additional_company in header_companies[1:]:
+                merged_output.append({
+                    "page_number": entry["page_number"],
+                    "pdf_name": entry["pdf_name"],
+                    "header": entry["header"],
+                    "label_name": entry["label_name"],
+                    "content": entry.get("content", ""),
+                    "table_content": entry.get("table_content", []),
+                    "matched_company": additional_company
+                })
+                
+        elif current_chunk:
+            # Continue adding to current chunk if no new company detected
+            if "content" in entry:
+                if current_chunk["content"]:
+                    current_chunk["content"] += "\n\n" + entry["content"]
+                    current_chunk["page_number"] += "," + str(entry["page_number"])
+                    page_numbers_list = list(dict.fromkeys(str(current_chunk["page_number"]).split(",")))
+                    page_numbers_list = [num.strip() for num in page_numbers_list if num.strip()]
+                    current_chunk["page_number"] = ",".join(page_numbers_list)
+
+                else:
+                    current_chunk["content"] = entry["content"]
+                    current_chunk["page_number"] = str(entry["page_number"])
+            
+            if "table_content" in entry:
+                current_chunk["table_content"].extend(entry["table_content"])
+                if current_chunk["page_number"]:
+                    if "metadata" in entry["table_content"]:
+                        if "table_page_id" in entry["table_content"]["metadata"]:
+                            current_chunk["page_number"] += "," + str(entry["table_content"]["metadata"]["table_page_id"])
+
+                    current_chunk["page_number"] += "," + str(entry["page_number"])
+                    page_numbers_list = list(dict.fromkeys(str(current_chunk["page_number"]).split(",")))
+                    page_numbers_list = [num.strip() for num in page_numbers_list if num.strip()]
+                    current_chunk["page_number"] = ",".join(page_numbers_list)
+
+            # if "page_number" in entry:
+            #     if current_chunk["page_number"]:
+            #         current_chunk["page_number"] += "," + str(entry["page_number"])
+            #     else:
+            #         current_chunk["page_number"] = str(entry["page_number"])
+
+        else:
+            # Ensure Unique page numbers for this entry
+            entry_copy = entry.copy()
+            if "page_number" in entry_copy :
+                page_numbers_list = list(dict.fromkeys(str(entry_copy["page_number"]).split(",")))
+                page_numbers_list = [num.strip() for num in page_numbers_list if num.strip()]
+                entry_copy["page_number"] = ",".join(page_numbers_list)
+
+            # Content before any company section
+            merged_output.append(entry_copy)
+    
+    # Add the last active chunk if it exists
+    if current_chunk:
+        # Ensure Unique page numbers for last entry
+        page_numbers_list = list(dict.fromkeys(str(current_chunk["page_number"]).split(",")))
+        page_numbers_list = [num.strip() for num in page_numbers_list if num.strip()]
+        entry_copy["page_number"] = ",".join(page_numbers_list)
+        merged_output.append(current_chunk)
+
+    merged_output_new = process_table_page_ids(merged_output=merged_output)
+    
+    return merged_output_new,fuzzy_matched_companies, portfolio_companies
+
+################################################################################################
+
+## Below code for using abbreviation funcnality
+
+import re
+
+def match_company_names(header_text: str, companies: List[str], threshold: int = FUZZY_MATCH_THRESHOLD) -> List[str]:
+    """Match company names in text, first checking header text abbreviations, then company abbreviations."""
+    header_text = str(header_text).lower().strip()
+    matched_companies = []
+    fuzzy_matched_companies = []
+    
+    # Generate possible abbreviations for header_text
+    header_abbreviations = [
+        ''.join(word[0] for word in header_text.split() if word),  # First letters of each word
+        re.sub(r'[aeiou\s]', '', header_text),  # Remove vowels and spaces
+        header_text.replace(' ', '')  # Remove spaces
+    ]
+    
+    for company in companies:
+        company_lower = company.lower()
+        
+        # First check: header text (full or abbreviated) against company full name
+        for header_pattern in [header_text] + header_abbreviations:
+            if fuzz.partial_ratio(header_pattern, company_lower) >= threshold:
+                matched_companies.append(company)
+                fuzzy_matched_companies.append(company)  # Record as fuzzy match
+                break
+        else:
+            # Second check: header text against company abbreviations
+            company_abbreviations = [
+                ''.join(word[0] for word in company_lower.split() if word),  # First letters of each word
+                re.sub(r'[aeiou\s]', '', company_lower),  # Remove vowels and spaces
+                company_lower.replace(' ', '')  # Remove spaces
+            ]
+            for company_pattern in company_abbreviations:
+                if fuzz.partial_ratio(header_text, company_pattern) >= threshold:
+                    matched_companies.append(company)
+                    fuzzy_matched_companies.append(company)  # Record as fuzzy match
+                    break
+
+    # Remove duplicates while preserving order
+    matched_companies = list(dict.fromkeys(matched_companies)) # Remove duplicates while preserving order
+    fuzzy_matched_companies = list(dict.fromkeys(fuzzy_matched_companies))   
+
+    return matched_companies, fuzzy_matched_companies
+
+
+################################################################################################################
+
+def process_document_company_wise(
+    intermediate_str_chunk_json: List[Dict],
+    output_directory: str,
+    file_name: str
+) -> List[Dict]:
+    """Process the document and return merged content in original format."""
+    # Convert string input to dict if needed
+    if isinstance(intermediate_str_chunk_json, str):
+        intermediate_str_chunk_json = json.loads(intermediate_str_chunk_json)
+    
+    # Merge content by company sections
+    # merged_content,matched_company_list = merge_portfolio_company_sections(intermediate_str_chunk_json)
+    merged_content,matched_company_list,portfolio_company_list = merge_portfolio_company_sections(intermediate_str_chunk_json)
+    # merged_content[0]["companies_list"] = matched_company_list
+    merged_content[0]["portfolio_companies_list_fuzzy_matched"] = matched_company_list
+    merged_content[0]["portfolio_companies_list_before"] = portfolio_company_list
+
+    print("matched_company_list::",matched_company_list)
+    print("portfolio_company_list::",portfolio_company_list)
+    
+    # Ensure output directory exists
+    os.makedirs(output_directory, exist_ok=True)
+    
+    # Save output
+    output_path = os.path.join(output_directory, f"{file_name}_h2h_merged_output.json")
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(merged_content, f, indent=4, ensure_ascii=False)
+    print(f"Saved merged output to {output_path}")
+    
+    return merged_content
+
+
+def read_json(file_path):
+    """Reads a JSON file and returns the parsed data."""
+    with open(file_path, 'r', encoding='utf-8') as file:
+        data = json.load(file)
+    return data
+
+
+# # Example usage
+if __name__ == "__main__":
+    input_str_chunk_json_path="/shared_disk/kushal/db_str_chunking/new_ws_structured_code/Triton2023Q4_patria_sample_output/Triton2023Q4_patria_sample_json_output/Triton2023Q4_patria_sample_final_h2h_extraction.json"
+    input_json = read_json(input_str_chunk_json_path)
+    
+    # Process the data
+    result = process_document_company_wise(
+        intermediate_str_chunk_json=input_json,
+        output_directory="db_structured_chunking/structure_chunking/src/iqeq_modification/testing_sample/output",
+        file_name="sample_report"
+    )
+    
+    print("Processing complete.")
+    # print(json.dumps(result, indent=2))
+