data_extractor.py

from transformers import AutoProcessor, AutoModelForTokenClassification

# Load the LayoutLM processor and model
processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base")
model = AutoModelForTokenClassification.from_pretrained("microsoft/layoutlmv3-base")

def extract_po_data_with_layoutlm(text):
    """
    Extract structured PO data using LayoutLM.

    Parameters:
    - text: Extracted plain text from the document.

    Returns:
    - Dictionary with structured PO data.
    """
    # Tokenize the input text
    inputs = processor(text, return_tensors="pt", padding=True, truncation=True)

    # Run the text through the LayoutLM model
    outputs = model(**inputs)
    predictions = outputs.logits.argmax(-1)

    # Convert predictions to labels
    tokens = processor.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    labels = [model.config.id2label[prediction] for prediction in predictions[0]]

    # Extract relevant entities
    po_data = {"Customer PO Number": None, "Items": []}
    current_item = {}
    for token, label in zip(tokens, labels):
        if label == "B-PO_NUMBER":
            po_data["Customer PO Number"] = token
        elif label == "B-ITEM_NAME":
            current_item["Item Name"] = token
        elif label == "B-QUANTITY":
            current_item["Quantity"] = token
        elif label == "B-RATE":
            current_item["Rate per Unit"] = token
        elif label == "B-UNIT":
            current_item["Unit of Measurement"] = token
        elif label == "B-DELIVERY_DATE":
            current_item["Delivery Date"] = token

        # End of an item entity
        if label == "O" and current_item:
            po_data["Items"].append(current_item)
            current_item = {}

    return po_data