Upload app.py

app.py

@@ -0,0 +1,549 @@
+import os
+from dotenv import load_dotenv
+
+import json
+import pandas as pd
+
+import zipfile
+import xml.etree.ElementTree as ET
+from io import BytesIO
+
+from openai import OpenAI
+
+import re
+
+import logging
+
+load_dotenv()
+HF_API_KEY = os.getenv("HF_API_KEY")
+
+# Configure logging to write to 'zaoju_logs.log' without using pickle
+logging.basicConfig(
+    filename='extract_po_logs.log',
+    level=logging.INFO, 
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    encoding='utf-8'
+)
+
+# Default Word XML namespace
+DEFAULT_NS = {'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'}
+NS = None  # Global variable to store the namespace
+
+def get_namespace(root):
+    """Extracts the primary namespace from the XML root element while keeping the default."""
+    global NS
+
+    if NS is None:
+        ns = root.tag.split('}')[0].strip('{')
+        NS = {'w': ns} if ns else DEFAULT_NS
+    return NS
+
+# --- Helper Functions for DOCX Processing ---
+
+def extract_text_from_cell(cell):
+    """Extracts text from a Word table cell, preserving line breaks and reconstructing split words."""
+    paragraphs = cell.findall('.//w:p', NS)
+    lines = []
+    
+    for paragraph in paragraphs:
+        # Get all text runs and concatenate their contents
+        text_runs = [t.text for t in paragraph.findall('.//w:t', NS) if t.text]
+        line = ''.join(text_runs).strip()  # Merge split words properly
+        
+        if line:  # Add only non-empty lines
+            lines.append(line)
+    
+    return lines  # Return list of lines to preserve line breaks
+
+def clean_spaces(text):
+    """
+    Removes excessive spaces between Chinese characters while preserving spaces in English words.
+    """
+    # Remove spaces **between** Chinese characters but keep English spaces
+    text = re.sub(r'([\u4e00-\u9fff])\s+([\u4e00-\u9fff])', r'\1\2', text)
+    return text.strip()
+
+def extract_key_value_pairs(text, target_dict=None):
+    """
+    Extracts multiple key-value pairs from a given text.
+    - First, split by more than 3 spaces (`\s{3,}`) **only if the next segment contains a `:`.**
+    - Then, process each segment by splitting at `:` to correctly assign keys and values.
+    """
+    if target_dict is None:
+        target_dict = {}
+
+    text = text.replace("：", ":")  # Normalize Chinese colons to English colons
+
+    # Step 1: Check if splitting by more than 3 spaces is necessary
+    segments = re.split(r'(\s{3,})', text)  # Use raw string to prevent invalid escape sequence
+
+    # Step 2: Process each segment, ensuring we only split if the next part has a `:`
+    merged_segments = []
+    temp_segment = ""
+
+    for segment in segments:
+        if ":" in segment:  # If segment contains `:`, it's a valid split point
+            if temp_segment:
+                merged_segments.append(temp_segment.strip())
+                temp_segment = ""
+            merged_segments.append(segment.strip())
+        else:
+            temp_segment += " " + segment.strip()
+
+    if temp_segment:
+        merged_segments.append(temp_segment.strip())
+
+    # Step 3: Extract key-value pairs correctly
+    for segment in merged_segments:
+        if ':' in segment:
+            key, value = segment.split(':', 1)  # Only split at the first colon
+            key, value = key.strip(), value.strip()  # Clean spaces
+
+            if key in target_dict:
+                target_dict[key] += "\n" + value  # Append if key already exists
+            else:
+                target_dict[key] = value
+
+    return target_dict
+
+# --- Table Processing Functions ---
+
+def process_single_column_table(rows):
+    """Processes a single-column table and returns the extracted lines as a list."""
+    single_column_data = []
+
+    for row in rows:
+        cells = row.findall('.//w:tc', NS)
+        if len(cells) == 1:
+            cell_lines = extract_text_from_cell(cells[0])  # Extract all lines from the cell
+            
+            # Append each line directly to the list without splitting
+            single_column_data.extend(cell_lines)
+
+    return single_column_data  # Return the list of extracted lines
+
+def process_buyer_seller_table(rows):
+    """Processes a two-column buyer-seller table into a structured dictionary using the first row as keys."""
+    headers = [extract_text_from_cell(cell) for cell in rows[0].findall('.//w:tc', NS)]
+    if len(headers) != 2:
+        return None  # Not a buyer-seller table
+
+    # determine role based on header text
+    def get_role(header_text, default_role):
+        header_text = header_text.lower()  # Convert to lowercase
+        if '买方' in header_text or 'buyer' in header_text or '甲方' in header_text:
+            return 'buyer_info'
+        elif '卖方' in header_text or 'seller' in header_text or '乙方' in header_text:
+            return 'seller_info'
+        else:
+            return default_role  # Default if no keyword is found
+
+    # Determine the keys for buyer and seller columns
+    buyer_key = get_role(headers[0][0], 'buyer_info')
+    seller_key = get_role(headers[1][0], 'seller_info')
+
+    # Initialize the dictionary using the determined keys
+    buyer_seller_data = {
+        buyer_key: {},
+        seller_key: {}
+    }
+
+    for row in rows:
+        cells = row.findall('.//w:tc', NS)
+        if len(cells) == 2:
+            buyer_lines = extract_text_from_cell(cells[0])
+            seller_lines = extract_text_from_cell(cells[1])
+
+        for line in buyer_lines:
+            extract_key_value_pairs(line, buyer_seller_data[buyer_key])
+
+        for line in seller_lines:
+            extract_key_value_pairs(line, buyer_seller_data[seller_key])
+            
+    return buyer_seller_data
+
+def process_summary_table(rows):
+    """Processes a two-column summary table where keys are extracted as dictionary keys."""
+    extracted_data = []
+    
+    for row in rows:
+        cells = row.findall('.//w:tc', NS)
+        if len(cells) == 2:
+            key = " ".join(extract_text_from_cell(cells[0]))
+            value = " ".join(extract_text_from_cell(cells[1]))
+            extracted_data.append({key: value})
+    
+    return extracted_data
+
+def extract_headers(first_row_cells):
+    """Extracts unique column headers from the first row of a table."""
+    headers = []
+    header_count = {}
+    for cell in first_row_cells:
+        cell_text = " ".join(extract_text_from_cell(cell))
+        grid_span = cell.find('.//w:gridSpan', NS)
+        col_span = int(grid_span.attrib.get(f'{{{NS["w"]}}}val', '1')) if grid_span is not None else 1
+        for _ in range(col_span):
+            # Ensure header uniqueness by appending an index if repeated
+            if cell_text in header_count:
+                header_count[cell_text] += 1
+                unique_header = f"{cell_text}_{header_count[cell_text]}"
+            else:
+                header_count[cell_text] = 1
+                unique_header = cell_text
+            headers.append(unique_header if unique_header else f"Column_{len(headers) + 1}")
+    return headers
+
+def process_long_table(rows):
+    """Processes a standard table and correctly handles horizontally merged cells."""
+    if not rows:
+        return []  # Avoid IndexError
+    
+    headers = extract_headers(rows[0].findall('.//w:tc', NS))
+    table_data = []
+    vertical_merge_tracker = {}
+
+    for row in rows[1:]:
+        row_data = {}
+        cells = row.findall('.//w:tc', NS)
+        running_index = 0
+        
+        for cell in cells:
+            cell_text = " ".join(extract_text_from_cell(cell))
+            
+            # Consistent Namespace Handling for Horizontal Merge
+            grid_span = cell.find('.//w:gridSpan', NS)
+            grid_span_val = grid_span.attrib.get(f'{{{NS["w"]}}}val') if grid_span is not None else '1'
+            col_span = int(grid_span_val)
+            
+            # Handle vertical merge
+            v_merge = cell.find('.//w:vMerge', NS)
+            if v_merge is not None:
+                v_merge_val = v_merge.attrib.get(f'{{{NS["w"]}}}val')
+                if v_merge_val == 'restart':
+                    vertical_merge_tracker[running_index] = cell_text
+                else:
+                    # Repeat the value from the previous row's merged cell
+                    cell_text = vertical_merge_tracker.get(running_index, "")
+            
+            # Repeat the value for horizontally merged cells
+            start_col = running_index
+            end_col = running_index + col_span
+            
+            # Repeat the value for each spanned column
+            for col in range(start_col, end_col):
+                key = headers[col] if col < len(headers) else f"Column_{col+1}"
+                row_data[key] = cell_text
+            
+            # Update the running index to the end of the merged cell
+            running_index = end_col
+        
+        # Fill remaining columns with empty strings to maintain alignment
+        while running_index < len(headers):
+            row_data[headers[running_index]] = ""
+            running_index += 1
+        
+        table_data.append(row_data)
+    
+    return table_data
+
+def extract_tables(root):
+    """Extracts tables from the DOCX document and returns structured data."""
+    tables = root.findall('.//w:tbl', NS)
+    table_data = {}
+    table_paragraphs = set()
+    
+    for table_index, table in enumerate(tables, start=1):
+        rows = table.findall('.//w:tr', NS)
+        if not rows:
+            continue  # Skip empty tables
+        
+        for paragraph in table.findall('.//w:p', NS):
+            table_paragraphs.add(paragraph)
+        
+        first_row_cells = rows[0].findall('.//w:tc', NS)
+        num_columns = len(first_row_cells)
+
+        if num_columns == 1:
+            single_column_data = process_single_column_table(rows)
+            if single_column_data:
+                table_data[f"table_{table_index}_single_column"] = single_column_data
+            continue  # Skip further processing for this table
+
+        summary_start_index = None
+        for i, row in enumerate(rows):
+            if len(row.findall('.//w:tc', NS)) == 2:
+                summary_start_index = i
+                break
+
+        long_table_data = []
+        summary_data = []
+
+        if summary_start_index is not None and summary_start_index > 0:
+            long_table_data = process_long_table(rows[:summary_start_index])
+        elif summary_start_index is None:
+            long_table_data = process_long_table(rows)
+
+        if summary_start_index is not None:
+            is_buyer_seller_table = all(len(row.findall('.//w:tc', NS)) == 2 for row in rows)
+            if is_buyer_seller_table:
+                buyer_seller_data = process_buyer_seller_table(rows)
+                if buyer_seller_data:
+                    table_data[f"table_{table_index}_buyer_seller"] = buyer_seller_data
+            else:
+                summary_data = process_summary_table(rows[summary_start_index:])
+    
+        if long_table_data:
+            table_data[f"long_table_{table_index}"] = long_table_data
+        if summary_data:
+            table_data[f"long_table_{table_index}_summary"] = summary_data
+    
+    return table_data, table_paragraphs
+
+# --- Non-Table Processing Functions ---
+
+def extract_text_outside_tables(root, table_paragraphs):
+    """Extracts text from paragraphs outside tables in the document."""
+    extracted_text = []
+    
+    # print(ET.tostring(root, encoding='unicode'))
+    for paragraph in root.findall('.//w:p', NS):
+        if paragraph in table_paragraphs:
+            continue  # Skip paragraphs inside tables
+        
+        texts = [t.text.strip() for t in paragraph.findall('.//w:t', NS) if t.text]
+        line = clean_spaces(' '.join(texts).replace(';', '').replace('；','').replace('：',':'))  # Remove semicolons and clean spaces
+        
+        if ':' in line:
+            extracted_text.append(line)
+    
+    return extracted_text
+
+# --- Main Extraction Functions ---
+
+def extract_docx_as_xml(file_bytes, save_xml=False, xml_filename="document.xml"):
+
+    # Ensure file_bytes is at the start position
+    file_bytes.seek(0)
+
+    with zipfile.ZipFile(file_bytes, 'r') as docx:
+        with docx.open('word/document.xml') as xml_file:
+            xml_content = xml_file.read().decode('utf-8')
+            if save_xml:
+                with open(xml_filename, "w", encoding="utf-8") as f:
+                    f.write(xml_content)
+            return xml_content
+
+def xml_to_json(xml_content, save_json=False, json_filename="extracted_data.json"):
+    
+    tree = ET.ElementTree(ET.fromstring(xml_content))
+    root = tree.getroot()
+    
+    table_data, table_paragraphs = extract_tables(root)
+    extracted_data = table_data
+    extracted_data["non_table_data"] = extract_text_outside_tables(root, table_paragraphs)
+    
+    if save_json:
+        with open(json_filename, "w", encoding="utf-8") as f:
+            json.dump(extracted_data, f, ensure_ascii=False, indent=4)
+    
+    return json.dumps(extracted_data, ensure_ascii=False, indent=4)
+
+def deepseek_extract_contract_summary(json_data, save_json=False):
+    """Sends extracted JSON data to OpenAI and returns formatted structured JSON."""
+    
+    # Step 1: Convert JSON string to Python dictionary
+    contract_data = json.loads(json_data)
+
+    # Step 2: Remove keys that contain "long_table"
+    filtered_contract_data = {key: value for key, value in contract_data.items() if "long_table" not in key}
+
+    # Step 3: Convert back to JSON string (if needed)
+    json_output = json.dumps(filtered_contract_data, ensure_ascii=False, indent=4)
+    
+    prompt = """You are given a contract in JSON format. Extract the following information:
+
+# Response Format
+Return the extracted information as a structured JSON in the exact format shown below (Do not repeat any keys):
+
+{
+    "合同编号":
+    "采购经办人":
+    "接收人":
+    "Recipient":
+    "接收地":
+    "Place of receipt":
+    "供应商":
+    "币种":
+    "合同日期":
+    "供货日期": 
+}
+
+Contract data in JSON format:""" + f"""
+{json_output}"""
+
+    messages = [
+        {
+            "role": "user",
+            "content": prompt
+        }
+    ]
+    
+    client = OpenAI(
+        base_url="https://router.huggingface.co/novita",
+        api_key=HF_API_KEY,
+    )
+    
+    completion = client.chat.completions.create(
+        model="deepseek/deepseek-r1-distill-qwen-14b", 
+        messages=messages, 
+    )
+    
+    contract_summary = re.sub(r"<think>.*?</think>\s*", "", completion.choices[0].message.content, flags=re.DOTALL) # Remove think
+    
+    contract_summary = re.sub(r"^```json\n|```$", "", contract_summary, flags=re.DOTALL) # Remove ```
+    
+    if save_json:
+        with open("extracted_contract_summary.json", "w", encoding="utf-8") as f:
+            f.write(contract_summary)
+    
+    return json.dumps(contract_summary, ensure_ascii=False, indent=4)
+
+def deepseek_extract_price_list(json_data):
+    """Sends extracted JSON data to OpenAI and returns formatted structured JSON."""
+    
+    # Step 1: Convert JSON string to Python dictionary
+    contract_data = json.loads(json_data)
+
+    # Step 2: Remove keys that contain "long_table"
+    filtered_contract_data = {key: value for key, value in contract_data.items() if "long_table" in key}
+
+    # Step 3: Convert back to JSON string (if needed)
+    json_output = json.dumps(filtered_contract_data, ensure_ascii=False, indent=4)
+    
+    print(json_output)
+    
+    prompt = """You are given a price list in JSON format. Extract the following information in CSV format:
+
+# Response Format
+Return the extracted information as a CSV in the exact format shown below:
+
+物料名称, 物料名称(英文), 物料规格, 采购数量, 单位, 单价, 计划号
+
+JSON data:""" + f"""
+{json_output}"""
+
+    messages = [
+        {
+            "role": "user",
+            "content": prompt
+        }
+    ]
+
+    client = OpenAI(
+	    base_url="https://router.huggingface.co/novita",
+	    api_key=HF_API_KEY,
+    )
+    
+    completion = client.chat.completions.create(
+        model="deepseek/deepseek-r1-distill-qwen-14b", 
+        messages=messages, 
+    )
+    
+    price_list = re.sub(r"<think>.*?</think>\s*", "", completion.choices[0].message.content, flags=re.DOTALL)
+    
+    price_list = re.sub(r"^```json\n|```$", "", price_list, flags=re.DOTALL)
+    
+    print(price_list)
+
+def json_to_excel(contract_summary, json_data, excel_path):
+    """Converts extracted JSON tables to an Excel file."""
+
+    # Correctly parse the JSON string
+    contract_summary_json = json.loads(json.loads(contract_summary))
+    
+    print(contract_summary_json)
+    
+    contract_summary_df = pd.DataFrame([contract_summary_json])
+    
+    # Ensure json_data is a dictionary
+    if isinstance(json_data, str):
+        json_data = json.loads(json_data)
+    
+    long_tables = [pd.DataFrame(table) for key, table in json_data.items() if "long_table" in key and "summary" not in key]
+    long_table = long_tables[-1] if long_tables else pd.DataFrame()
+
+    with pd.ExcelWriter(excel_path) as writer:
+        contract_summary_df.to_excel(writer, sheet_name="Contract Summary", index=False)
+        long_table.to_excel(writer, sheet_name="Price List", index=False)
+
+#--- Extract PO ------------------------------
+
+def extract_po(docx_path):
+    """Processes a single .docx file, extracts tables, formats with OpenAI, and saves as an Excel file."""
+    if not os.path.exists(docx_path) or not docx_path.endswith(".docx"):
+        print(f"Invalid file: {docx_path}")
+        return
+    
+    # Read the .docx file as bytes
+    with open(docx_path, "rb") as f:
+        docx_bytes = BytesIO(f.read())
+        
+    # Step 1: Extract XML content from DOCX
+    print("Extracting Docs data to XML...")
+    xml_file = extract_docx_as_xml(docx_bytes,save_xml=True)
+    
+    get_namespace(ET.fromstring(xml_file))
+    
+    # Step 2: Extract tables from DOCX and save JSON
+    print("Extracting XML data to JSON...")
+    extracted_data = xml_to_json(xml_file, save_json=True)
+
+    # Step 2: Process JSON with OpenAI to get structured output
+    print("Processing JSON data with AI...")
+    contract_summary = deepseek_extract_contract_summary(extracted_data, save_json=True)
+
+    # Step 3: Save formatted data as Excel
+    print("Converting AI Generated JSON to Excel...")
+    excel_output_path = os.path.splitext(docx_path)[0] + ".xlsx"
+    json_to_excel(contract_summary, extracted_data, excel_output_path)
+
+    print(f"Excel file saved at: {excel_output_path}")
+    
+    
+    # Logging
+    log = f"""Results:
+    
+    Contract Summary: {contract_summary},
+    
+    RAW Extracted Data: {extracted_data},
+    
+    XML Preview: {xml_file[:1000]}"""
+    
+    print(log)
+    
+    logging.info(f"""{log}""")
+    
+    
+    return excel_output_path
+
+# Example Usage
+
+# extract_po("test-contract-converted.docx")
+# extract_po("test-contract.docx")
+
+# Gradio Interface ------------------------------
+
+import gradio as gr
+from gradio.themes.base import Base
+
+interface = gr.Interface(
+    fn=extract_po,
+    title="PO Extractor 买卖合同数据提取", 
+    inputs=gr.File(label="买卖合同 （.docx）"),
+    outputs=gr.File(label="数据提取结果 （.xlsx）"),
+    allow_flagging="never",
+    theme=Base()
+)
+
+interface.launch()