data-science-document-ai 1.38.0__py3-none-any.whl → 1.39.0__py3-none-any.whl

{data_science_document_ai-1.38.0.dist-info → data_science_document_ai-1.39.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-science-document-ai
-Version: 1.38.0
+Version: 1.39.0
 Summary: "Document AI repo for data science"
 Author: Naomi Nguyen
 Author-email: naomi.nguyen@forto.com

{data_science_document_ai-1.38.0.dist-info → data_science_document_ai-1.39.0.dist-info}/RECORD

@@ -6,11 +6,11 @@ src/excel_processing.py,sha256=ZUlZ5zgWObmQfAWHoSrEEITKwr-xXxuOiPC3qDnGjtQ,2459
 src/io.py,sha256=IXz4wWqiHa9mnHNgtrC6X9M2lItYp9eu6rHCThUIh5c,3585
 src/llm.py,sha256=aEK3rL8XvY7CakvkOJQmcHpEKwZRd8PPrLrzHiO-GFk,7827
 src/log_setup.py,sha256=RhHnpXqcl-ii4EJzRt47CF2R-Q3YPF68tepg_Kg7tkw,2895
-src/pdf_processing.py,sha256=g-WVrI6J2lbVR3eOoDJnuy4buWh7bTmO-3aezoTN3i4,15527
-src/postprocessing/common.py,sha256=UxwmnXH7saggxDMs9Ssx_Bp3-O9NeUrcKFWRI_QYuZ0,39583
+src/pdf_processing.py,sha256=yB0FpIdSRqxeEbZAIK_bPFypWrSSMb8uwCRxTTFfmxc,15493
+src/postprocessing/common.py,sha256=OR9O73gUP4tevIZMnorbiUgzviEJlVr46ArTWMXrYVA,19316
 src/postprocessing/postprocess_booking_confirmation.py,sha256=nK32eDiBNbauyQz0oCa9eraysku8aqzrcoRFoWVumDU,4827
 src/postprocessing/postprocess_commercial_invoice.py,sha256=3I8ijluTZcOs_sMnFZxfkAPle0UFQ239EMuvZfDZVPg,1028
-src/postprocessing/postprocess_partner_invoice.py,sha256=oCT-l31DTosUf0cz0d5IWOF6erw6rD3rQfR58koSUeM,11760
+src/postprocessing/postprocess_partner_invoice.py,sha256=bWm3Miaq_mtX62xSs14vNQCWPHOj2895Bt6TuOVZWZU,11742
 src/prompts/library/bookingConfirmation/evergreen/placeholders.json,sha256=Re2wBgZoaJ5yImUUAwZOZxFcKXHxi83TCZwTuqd2v2k,1405
 src/prompts/library/bookingConfirmation/evergreen/prompt.txt,sha256=qlBMFDHy-gwr2PVeuHrfMEg_8Ibdym243DnaCgINa7g,2614
 src/prompts/library/bookingConfirmation/hapag-lloyd/placeholders.json,sha256=Re2wBgZoaJ5yImUUAwZOZxFcKXHxi83TCZwTuqd2v2k,1405
@@ -51,9 +51,9 @@ src/prompts/library/preprocessing/carrier/placeholders.json,sha256=1UmrQNqBEsjLI
 src/prompts/library/preprocessing/carrier/prompt.txt,sha256=NLvRZQCZ6aWC1yTr7Q93jK5z7Vi_b4HBaiFYYnIsO-w,134
 src/prompts/library/shippingInstruction/other/prompt.txt,sha256=fyC24ig4FyRNnLuQM69s4ZVajsK-LHIl2dvaaEXr-6Q,1327
 src/prompts/prompt_library.py,sha256=VJWHeXN-s501C2GiidIIvQQuZdU6T1R27hE2dKBiI40,2555
-src/setup.py,sha256=TJu68mXS6Dx90Il8A_pHDnrIOiLD3q9f7FWgW0c1HOM,8352
+src/setup.py,sha256=kPSZosrICfaGZeDaajr40Ha7Ok4XK4fo_uq35Omiwr0,7128
 src/tms.py,sha256=UXbIo1QE--hIX6NZi5Qyp2R_CP338syrY9pCTPrfgnE,1741
-src/utils.py,sha256=68x3hakQ8aDfq7967XoTRe_vsneWnLbWp_jz8q_FrBA,12189
-data_science_document_ai-1.38.0.dist-info/METADATA,sha256=fDQUs1Zi2ZteWsX6vm-0pNEZA04KJAels_OjC4Job6o,2153
-data_science_document_ai-1.38.0.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
-data_science_document_ai-1.38.0.dist-info/RECORD,,
+src/utils.py,sha256=fTzXSeXejSmeGH0thOEdtDTAVhmeALVg44ORFarBKOk,13826
+data_science_document_ai-1.39.0.dist-info/METADATA,sha256=odM1ywura_FepOjD6iVQXc4QXVsmrtRQILScEsyZTM0,2153
+data_science_document_ai-1.39.0.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
+data_science_document_ai-1.39.0.dist-info/RECORD,,

src/pdf_processing.py

@@ -355,7 +355,6 @@ async def data_extraction_manual_flow(
     meta,
     processor_client,
     schema_client,
-    embed_manager,
 ):
     """
     Process a PDF file and extract data from it.
@@ -418,7 +417,7 @@ async def data_extraction_manual_flow(
         )
     # Create the result dictionary with the extracted data
     extracted_data = await format_all_entities(
-        extracted_data, embed_manager, meta.documentTypeCode, params
+        extracted_data, meta.documentTypeCode, params
     )
     result = {
         "id": meta.id,

src/postprocessing/common.py

@@ -1,426 +1,22 @@
 import asyncio
-import datetime
 import json
 import os
 import re
 from datetime import timezone
 
-import numpy as np
 import pandas as pd
-import requests
 from nltk.corpus import stopwords
 from rapidfuzz import process
-from vertexai.preview.language_models import TextEmbeddingModel
 
 from src.constants import formatting_rules
-from src.io import get_storage_client, logger
+from src.io import logger
 from src.postprocessing.postprocess_partner_invoice import process_partner_invoice
 from src.prompts.prompt_library import prompt_library
-from src.tms import call_tms, set_tms_service_token
+from src.utils import get_tms_mappings
 
 tms_domain = os.environ["TMS_DOMAIN"]
 
 
-class EmbeddingsManager:  # noqa: D101
-    def __init__(self, params):  # noqa: D107
-        self.params = params
-        self.embeddings_dict = {}
-        self.embed_model = setup_embed_model()
-        self.bucket = self.get_bucket_storage()
-        self.embedding_folder = self.embed_model._model_id
-        self.embedding_dimension = 768  # TODO: to be reduced
-
-    def get_bucket_storage(self):
-        """
-        Retrieve the bucket storage object.
-
-        Returns:
-            The bucket storage object.
-        """
-        params = self.params
-        storage_client = get_storage_client(params)
-        bucket = storage_client.bucket(params["doc_ai_bucket_name"])
-        return bucket
-
-    def _find_most_similar_option(self, input_string, option_ids, option_embeddings):
-        """
-        Find the most similar option to the given input string based on embeddings.
-
-        Args:
-            model: The model used for generating embeddings.
-            input_string (str): The input string to find the most similar option for.
-            option_ids (list): The list of option IDs.
-            option_embeddings (np.ndarray): The embeddings of the options.
-
-        Returns:
-            The ID of the most similar option.
-        """
-        try:
-            input_embedding = self.embed_model.get_embeddings(
-                [input_string], output_dimensionality=self.embedding_dimension
-            )[0].values
-            similarities = np.dot(option_embeddings, input_embedding)
-            idx = np.argmax(similarities)
-            return option_ids[idx]
-        except Exception as e:
-            logger.error(f"Embeddings error: {e}")
-            return None
-
-    def load_embeddings(self):
-        """
-        Load embeddings for container types, ports, and terminals.
-
-        Returns:
-            None
-        """
-        for data_field in [
-            "container_types",
-            "ports",
-            "terminals",
-            "depots",
-            "item_codes_label",
-        ]:
-            self.embeddings_dict[data_field] = load_embed_by_data_field(
-                self.bucket,
-                f"{self.embedding_folder}/{data_field}/output",
-                self.embedding_dimension,
-            )
-
-    async def update_embeddings(self):
-        """
-        Update the embeddings dictionary.
-
-        Returns:
-            dict: The updated embeddings dictionary with the following keys:
-                - "container_types": A tuple containing the container types and their embeddings.
-                - "ports": A tuple containing the ports and their embeddings.
-                - "terminals": A tuple containing the terminal IDs and their embeddings.
-        """
-        # Update embeddings dict here.
-        # Ensure this method is async if you're calling async operations.
-        set_tms_service_token()
-        (
-            container_types,
-            container_type_embeddings,
-        ) = self.setup_container_type_embeddings(
-            *self.embeddings_dict.get("container_types", ([], []))
-        )
-
-        ports, port_embeddings = self.setup_ports_embeddings(
-            *self.embeddings_dict.get("ports", ([], []))
-        )
-
-        # Setup terminal embeddings
-        # Since retrieving terminal attributes requires calling TMS' api to extract terminals by each port,
-        # we only do it for new ports.
-        prev_port_ids, _ = self.embeddings_dict.get("ports", ([], []))
-        added_port_ids = [port for port in ports if port not in prev_port_ids]
-        if added_port_ids:
-            terminal_ids, terminal_embeddings = self.setup_terminal_embeddings(
-                added_port_ids
-            )
-        else:
-            terminal_ids, terminal_embeddings = self.embeddings_dict["terminals"]
-
-        depot_names, depot_embeddings = self.setup_depot_embeddings(
-            *self.embeddings_dict.get("depots", ([], []))
-        )
-
-        item_code_names, item_code_embeddings = self.setup_item_code_embeddings(
-            *self.embeddings_dict.get("item_codes_label", ([], []))
-        )
-
-        self.embeddings_dict = {
-            "container_types": (container_types, container_type_embeddings),
-            "ports": (ports, port_embeddings),
-            "terminals": (terminal_ids, terminal_embeddings),
-            "depots": (depot_names, depot_embeddings),
-            "item_codes_label": (item_code_names, item_code_embeddings),
-        }
-        return self.embeddings_dict
-
-    def batch_embed(self, option_strings: list[dict], suffix: str):
-        """
-        Compute embeddings for a batch of option strings and uploads them to a cloud storage bucket.
-
-        Args:
-            option_strings (list): A list of option strings to compute embeddings for.
-            suffix (str): A suffix to be used in the storage path for the embeddings:
-            input & output will be stored under "{bucket}/{parent_folder}/{suffix}/"
-
-        Returns:
-            tuple: A tuple containing the option IDs and embeddings.
-        """
-        now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-        input_path = f"{self.embedding_folder}/{suffix}/input/{now}.jsonl"
-        blob = self.bucket.blob(input_path)
-
-        # Convert each dictionary to a JSON string and join them with newlines
-        option_strings = [
-            {**option, "task_type": "SEMANTIC_SIMILARITY", "output_dimensionality": 256}
-            for option in option_strings
-        ]
-        jsonl_string = "\n".join(json.dumps(d) for d in option_strings)
-
-        # Convert the combined string to bytes
-        jsonl_bytes = jsonl_string.encode("utf-8")
-
-        # Upload the bytes to the blob
-        blob.upload_from_string(jsonl_bytes, content_type="text/plain")
-
-        # Compute embeddings for the options
-        embedding_path = f"{self.embedding_folder}/{suffix}/output"
-        assert len(option_strings) <= 30000  # Limit for batch embedding
-        batch_resp = self.embed_model.batch_predict(
-            dataset=f"gs://{self.bucket.name}/{input_path}",  # noqa
-            destination_uri_prefix=f"gs://{self.bucket.name}/{embedding_path}",  # noqa
-        )
-
-        if batch_resp.state.name != "JOB_STATE_SUCCEEDED":
-            logger.warning(
-                f"Batch prediction job failed with state {batch_resp.state.name}"
-            )
-        else:
-            logger.info(f"Embeddings for {suffix} computed successfully.")
-
-        option_ids, option_embeddings = load_embed_by_data_field(
-            self.bucket, embedding_path, self.embedding_dimension
-        )
-        return option_ids, option_embeddings
-
-    def setup_container_type_embeddings(
-        self, computed_container_type_ids, computed_container_type_embeddings
-    ):
-        """
-        Set up container type embeddings.
-
-        Args:
-            computed_container_type_ids (list): The list of already computed container type IDs.
-            computed_container_type_embeddings (list): The list of already computed container type embeddings.
-
-        Returns:
-            tuple: A tuple containing the updated container type IDs and embeddings.
-        """
-        url = (
-            f"https://tms.forto.{tms_domain}/api/transport-units/api/types/list"  # noqa
-        )
-        resp = call_tms(requests.get, url)
-        container_types = resp.json()
-
-        container_attribute_strings = [
-            dict(
-                title=container_type["code"],
-                content=" | ".join(
-                    [container_type["code"]]
-                    + [
-                        f"{v}"
-                        for k, v in container_type["containerAttributes"].items()
-                        if k in ["isoSizeType", "isoTypeGroup", "containerCategory"]
-                    ]
-                    + [container_type.get(k, "") for k in ["displayName", "notes"]]
-                ),
-            )
-            for container_type in container_types
-            if container_type["type"] == "Container"
-            and container_type["code"] not in computed_container_type_ids
-            and container_type.get("containerAttributes") is not None
-        ]
-        if not container_attribute_strings:
-            logger.info("No new container types found.")
-            return computed_container_type_ids, computed_container_type_embeddings
-
-        logger.info("Computing embeddings for container types...")
-        container_type_ids, container_type_embeddings = self.batch_embed(
-            container_attribute_strings, "container_types"
-        )
-        return container_type_ids, container_type_embeddings
-
-    def setup_ports_embeddings(self, computed_port_ids, computed_port_embeddings):
-        """
-        Set up port embeddings.
-
-        Steps:
-        - Retrieve active ports from the TMS API
-        - Compute embeddings for new tradelane-enabled ports
-        - Return ALL port IDs and embeddings.
-
-        Args:
-            computed_port_ids (list): The list of previously computed port IDs.
-            computed_port_embeddings (list): The list of previously computed port embeddings.
-
-        Returns:
-            tuple: A tuple containing ALL port IDs and embeddings.
-        """
-        url = f"https://tms.forto.{tms_domain}/api/transport-network/api/ports?pageSize=1000000&status=active"  # noqa
-        resp = call_tms(requests.get, url)
-        resp_json = resp.json()
-        if len(resp_json["data"]) != resp_json["_paging"]["totalRecords"]:
-            logger.error("Not all ports were returned.")
-
-        new_sea_ports = [
-            port
-            for port in resp_json["data"]
-            if "sea" in port["modes"] and port["id"] not in computed_port_ids
-        ]
-        if not new_sea_ports:
-            logger.info("No new ports found.")
-            return computed_port_ids, computed_port_embeddings
-
-        port_attribute_strings = [
-            dict(
-                title=port["id"],
-                content=" ".join(
-                    [
-                        "port for shipping",
-                        add_text_without_space(
-                            port["name"]
-                        ),  # for cases like QUINHON - Quinhon
-                        port["id"],
-                    ]
-                ),
-            )
-            for port in new_sea_ports
-        ]
-
-        logger.info("Computing embeddings for ports.")
-        port_ids, port_embeddings = self.batch_embed(port_attribute_strings, "ports")
-        return port_ids, port_embeddings
-
-    def setup_depot_embeddings(self, computed_depot_names, computed_depot_embeddings):
-        """
-        Set up depot embeddings.
-
-        Steps:
-        - Retrieve active depot from the TMS API
-        - Compute embeddings for new tdepot
-        - Return ALL depot names and embeddings.
-
-        Args:
-            computed_depot_names (list): The list of previously computed depot names.
-            computed_depot_embeddings (list): The list of previously computed depot embeddings.
-
-        Returns:
-            tuple: A tuple containing ALL depot names and embeddings.
-        """
-        url = f"https://tms.forto.{tms_domain}/api/transport-network/api/depots?pageSize=1000000"  # noqa
-        resp = call_tms(requests.get, url)
-        resp_json = resp.json()
-
-        new_depots = [
-            depot
-            for depot in resp_json["data"]
-            if depot["name"] not in computed_depot_names
-        ]
-        if not new_depots:
-            logger.info("No new depots found.")
-            return computed_depot_names, computed_depot_embeddings
-
-        depot_attribute_strings = [
-            dict(
-                title=depot["name"],
-                content=" | ".join(
-                    [
-                        "depot",
-                        "name - " + depot["name"],
-                        "address - " + depot["address"]["fullAddress"],
-                    ]
-                ),
-            )
-            for depot in resp_json["data"]
-        ]
-
-        logger.info("Computing embeddings for depots.")
-        depot_names, depot_embeddings = self.batch_embed(
-            depot_attribute_strings, "depots"
-        )
-        return depot_names, depot_embeddings
-
-    def setup_terminal_embeddings(self, added_port_ids):
-        """
-        Set up terminal embeddings for `added_port_ids`, using `model`, uploaded to `bucket`.
-
-        Args:
-            added_port_ids (list): A list of added port IDs.
-
-        Returns:
-            tuple: A tuple containing the ALL terminal IDs and terminal embeddings.
-            Not just for the added port IDs.
-        """
-        terminal_attibute_strings = [
-            setup_terminal_attributes(port_id) for port_id in added_port_ids
-        ]
-        terminal_attibute_strings = sum(terminal_attibute_strings, [])
-        if not terminal_attibute_strings:
-            logger.info("No new terminals found.")
-            return [], np.array([])
-
-        terminal_ids, terminal_embeddings = self.batch_embed(
-            terminal_attibute_strings, "terminals"
-        )
-        return terminal_ids, terminal_embeddings
-
-    def setup_item_code_embeddings(
-        self, computed_item_code_names, computed_item_code_embeddings
-    ):
-        """
-        Set up item_code embeddings.
-
-        Steps:
-        - Retrieve active item_code from the TMS API
-        - Compute embeddings for new titem_code
-        - Return ALL item_code names and embeddings.
-
-        Args:
-            computed_item_code_names (list): The list of previously computed item_code names.
-            computed_item_code_embeddings (list): The list of previously computed item_code embeddings.
-
-        Returns:
-            tuple: A tuple containing ALL item_code names and embeddings.
-        """
-        url = f"https://tms.forto.{tms_domain}/api/catalog/item-codes?transportTypes=fcl&pageSize=1000000"  # noqa
-        resp = call_tms(requests.get, url)
-        resp_json = resp.json()
-
-        new_item_codes = [
-            item_code
-            for item_code in resp_json["results"]
-            if item_code["id"] not in computed_item_code_names
-        ]
-        if not new_item_codes:
-            logger.info("No new item_codes found.")
-            return computed_item_code_names, computed_item_code_embeddings
-
-        item_code_attribute_strings = [
-            dict(
-                title=item_code["id"],
-                content=" | ".join(
-                    [
-                        item_code["id"],
-                        item_code["label"],
-                    ]
-                ),
-            )
-            for item_code in resp_json["results"]
-        ]
-
-        logger.info("Computing embeddings for item_codes.")
-        item_code_names, item_code_embeddings = self.batch_embed(
-            item_code_attribute_strings, "item_codes_label"
-        )
-        return item_code_names, item_code_embeddings
-
-
-def setup_embed_model():
-    """
-    Set up and return a text embedding model.
-
-    Returns:
-        TextEmbeddingModel: The initialized text embedding model.
-    """
-    model = TextEmbeddingModel.from_pretrained("text-multilingual-embedding-002")
-    return model
-
-
 def convert_container_number(container_number):
     """
     Convert a container number to ISO standard.
@@ -570,106 +166,6 @@ def extract_string(data_field_value):
     return formatted_value if formatted_value not in ["''", ""] else None
 
 
-def extract_google_embed_resp(prediction_string, embedding_dimension):
-    """
-    Extract relevant information from the Google Embed API response.
-
-    Args:
-        prediction_string (str): The prediction string returned by the Google Embed API.
-
-    Returns:
-        dict: A dictionary containing the extracted information.
-            - _id (str): The title of the instance.
-            - attr_text (str): The content of the instance.
-            - embedding (list): The embeddings values from the predictions.
-
-    """
-    res = json.loads(prediction_string)
-    return dict(
-        _id=res["instance"]["title"],
-        attr_text=res["instance"]["content"],
-        embedding=res["predictions"][0]["embeddings"]["values"][:embedding_dimension],
-    )
-
-
-def load_embed_by_data_field(bucket, embedding_path, embedding_dimension):
-    """
-    Load embeddings by data field from the specified bucket and embedding path.
-
-    Args:
-        bucket (Bucket): The bucket object representing the storage bucket.
-        embedding_path (str): The path to the embeddings in the bucket (different by data_field).
-
-    Returns:
-        tuple: A tuple containing the option IDs and option embeddings.
-            - option_ids (list): A list of option IDs.
-            - option_embeddings (ndarray): An array of option embeddings.
-    """
-    # Retrieve the embeddings from the output files
-    blobs = bucket.list_blobs(prefix=embedding_path)
-    all_blob_data = []
-    for blob in blobs:
-        blob_data = blob.download_as_bytes().decode("utf-8").splitlines()
-        embeddings = [
-            extract_google_embed_resp(data, embedding_dimension) for data in blob_data
-        ]
-        all_blob_data.extend(embeddings)
-    option_ids = [embed["_id"] for embed in all_blob_data]
-    option_embeddings = np.stack([embed["embedding"] for embed in all_blob_data])
-    return option_ids, option_embeddings
-
-
-def setup_terminal_attributes(port_id: str):
-    """
-    Retrieve and format the attributes of active terminals at a given port.
-
-    Args:
-        port_id (str): The ID of the port.
-
-    Returns:
-        list: A list of dictionaries containing the formatted attributes of active terminals.
-              Each dictionary has the following keys:
-              - title: The terminal's short code.
-              - content: A string representation of the terminal's attributes, including its name,
-                         searchable name, and full address.
-    """
-    url = f"https://gateway.forto.{tms_domain}/api/transport-network/api/ports/{port_id}/terminals/list"  # noqa
-    resp = call_tms(requests.get, url)
-    terminals = resp.json()
-    if len(terminals) == 0:
-        return []
-    active_terminals = [term for term in terminals if term["isActive"]]
-    if len(active_terminals) == 0:
-        logger.warning(f"No active terminals found at port {port_id}.")
-        return []
-
-    terminal_attibute_strings = [
-        dict(
-            title=term["name"],
-            content=" | ".join(
-                [
-                    "shipping terminal",
-                    "code - " + term["terminalShortCode"],
-                    "name - " + modify_terminal_name(term["searchableName"]),
-                    "country - " + term["address"]["country"],
-                ]
-            ),
-        )
-        for term in active_terminals
-    ]
-    return terminal_attibute_strings
-
-
-def modify_terminal_name(text):
-    # Find the first occurrence of a word starting with 'K' followed by a number
-    # and replace it with 'KAAI' - meaning Quay in Dutch
-    match = re.search(r"K(\d+)", text)
-    if match:
-        # Append "KAAI" followed by the number if a match is found
-        text += f" KAAI {match.group(1)}"
-    return text
-
-
 def remove_none_values(d):
     if isinstance(d, dict):
         # Create a new dictionary to store non-None values
@@ -731,25 +227,6 @@ def convert_invoice_type(data_field_value, params):
     return None
 
 
-def validate_reverse_charge_value(reverse_charge_sentence_value):
-    """
-    Validates the reverseChargeSentence value before assigning to line items.
-
-    Args:
-        reverse_charge_sentence_value (bool): The formatted value of reverseChargeSentence (True or False).
-
-    Returns:
-        bool: The validated reverseChargeSentence value.
-    """
-    if isinstance(reverse_charge_sentence_value, bool):
-        return reverse_charge_sentence_value
-    else:
-        logger.warning(
-            f"Invalid reverseChargeSentence value: {reverse_charge_sentence_value}. Defaulting to False."
-        )
-        return False
-
-
 # Function to create KVP dictionary using apply method
 def create_kvp_dictionary(df_raw: pd.DataFrame):
     """Create a key-value pair dictionary from the given DataFrame.
@@ -887,47 +364,37 @@ def clean_item_description(lineitem: str, remove_numbers: bool = True):
     return re.sub(r"\s{2,}", " ", lineitem).strip()
 
 
-async def format_label(
-    entity_k, entity_value, embed_manager, document_type_code, params
-):
+async def format_label(entity_k, entity_value, document_type_code, params):
     llm_client = params["LlmClient"]
     if isinstance(entity_value, dict):  # if it's a nested entity
         format_tasks = [
-            format_label(sub_k, sub_v, embed_manager, document_type_code, params)
+            format_label(sub_k, sub_v, document_type_code, params)
             for sub_k, sub_v in entity_value.items()
         ]
         return entity_k, {k: v for k, v in await asyncio.gather(*format_tasks)}
     if isinstance(entity_value, list):
         format_tasks = await asyncio.gather(
             *[
-                format_label(entity_k, sub_v, embed_manager, document_type_code, params)
+                format_label(entity_k, sub_v, document_type_code, params)
                 for sub_v in entity_value
             ]
         )
         return entity_k, [v for _, v in format_tasks]
     entity_key = entity_k.lower()
-    embeddings_dict = embed_manager.embeddings_dict
     formatted_value = None
 
     if entity_key.startswith("port"):
-        formatted_value = await get_port_code_ai(
-            entity_value, llm_client, embed_manager, *embeddings_dict["ports"]
-        )
+        formatted_value = await get_port_code_ai(entity_value, llm_client)
+
     elif (entity_key == "containertype") or (entity_key == "containersize"):
-        formatted_value = embed_manager._find_most_similar_option(
-            "container type " + entity_value,
-            *embeddings_dict["container_types"],
-        )
+        formatted_value = get_tms_mappings(entity_value, "container_types")
+
     elif check_formatting_rule(entity_k, document_type_code, "terminal"):
-        formatted_value = embed_manager._find_most_similar_option(
-            "shipping terminal " + str(entity_value),
-            *embeddings_dict["terminals"],
-        )
+        formatted_value = get_tms_mappings(entity_value, "terminals")
+
     elif check_formatting_rule(entity_k, document_type_code, "depot"):
-        formatted_value = embed_manager._find_most_similar_option(
-            "depot " + str(entity_value),
-            *embeddings_dict["depots"],
-        )
+        formatted_value = get_tms_mappings(entity_value, "depots")
+
     elif entity_key.startswith(("eta", "etd", "duedate", "issuedate", "servicedate")):
         try:
             cleaned_data_field_value = clean_date_string(entity_value)
@@ -1002,15 +469,11 @@ async def format_label(
     return entity_k, result
 
 
-async def get_port_code_ai(
-    port: str, llm_client, embed_manager, port_ids, port_embeddings
-):
+async def get_port_code_ai(port: str, llm_client):
+    """Get port code using AI model."""
     port_llm = await get_port_code_llm(port, llm_client)
 
-    if port_llm in port_ids:
-        return port_llm
-    port_text = f"port for shipping {port}"
-    return embed_manager._find_most_similar_option(port_text, port_ids, port_embeddings)
+    return get_tms_mappings(port, "ports", port_llm)
 
 
 async def get_port_code_llm(port: str, llm_client):
@@ -1081,7 +544,7 @@ def decimal_convertor(value):
     return value
 
 
-async def format_all_entities(result, embed_manager, document_type_code, params):
+async def format_all_entities(result, document_type_code, params):
     """Format the entity values in the result dictionary."""
     # Since we treat `customsInvoice` same as `partnerInvoice`
     document_type_code = (
@@ -1096,15 +559,11 @@ async def format_all_entities(result, embed_manager, document_type_code, params)
         return {}
 
     # Format all entities recursively
-    _, aggregated_data = await format_label(
-        None, result, embed_manager, document_type_code, params
-    )
+    _, aggregated_data = await format_label(None, result, document_type_code, params)
 
     # Process partner invoice on lineitem mapping and reverse charge sentence
     if document_type_code in ["partnerInvoice", "bundeskasse"]:
-        process_partner_invoice(
-            params, aggregated_data, embed_manager, document_type_code
-        )
+        process_partner_invoice(params, aggregated_data, document_type_code)
 
     logger.info("Data Extraction completed successfully")
     return aggregated_data

src/postprocessing/postprocess_partner_invoice.py

@@ -4,6 +4,7 @@ from concurrent.futures import ThreadPoolExecutor
 from fuzzywuzzy import fuzz
 
 from src.io import logger
+from src.utils import get_tms_mappings
 
 
 def postprocessing_partner_invoice(partner_invoice):
@@ -135,7 +136,7 @@ def update_recipient_and_vendor(aggregated_data, is_recipient_forto):
     ] = "Dasbachstraße 15, 54292 Trier, Germany"
 
 
-def process_partner_invoice(params, aggregated_data, embed_manager, document_type_code):
+def process_partner_invoice(params, aggregated_data, document_type_code):
     """Process the partner invoice data."""
     # Post process containerNumber.
     # TODO: Remove this block of code after migrating to LLM completely and update the placeholder in the prompt library
@@ -192,7 +193,6 @@ def process_partner_invoice(params, aggregated_data, embed_manager, document_typ
     for line_item in line_items:
         if line_item.get("lineItemDescription", None) is not None:
             line_item["itemCode"] = associate_forto_item_code(
-                embed_manager,
                 line_item["lineItemDescription"]["formattedValue"],
                 params,
             )
@@ -275,7 +275,7 @@ def find_matching_lineitem(new_lineitem: str, kvp_dict: dict, threshold=90):
     return kvp_dict.get(best_match, None)
 
 
-def associate_forto_item_code(embed_manager, input_string, params):
+def associate_forto_item_code(input_string, params):
     """
     Finds a match for the input string using fuzzy matching first, then embedding fallback.
 
@@ -286,7 +286,6 @@ def associate_forto_item_code(embed_manager, input_string, params):
 
     Args:
         input_string: The string to find a match for.
-        embed_manager: The embedding manager instance to use for fallback.
         params: Parameters containing the lookup data and fuzzy threshold.
 
     Returns:
@@ -301,10 +300,11 @@ def associate_forto_item_code(embed_manager, input_string, params):
 
     if forto_item_code is None:
         # 2. Fallback to embedding function if no good fuzzy match
-        embeddings_dict = embed_manager.embeddings_dict
-        forto_item_code = embed_manager._find_most_similar_option(
-            input_string, *embeddings_dict["item_codes_label"]
-        )
+        forto_item_code = get_tms_mappings(input_string, "line_items")
+        # embeddings_dict = embed_manager.embeddings_dict
+        # forto_item_code = embed_manager._find_most_similar_option(
+        #     input_string, *embeddings_dict["item_codes_label"]
+        # )
 
     result = {"documentValue": input_string, "formattedValue": forto_item_code}
     return result

src/setup.py

@@ -1,5 +1,4 @@
 """Contains project setup parameters and initialization functions."""
-import argparse
 import json
 
 # import streamlit as st
@@ -69,50 +68,6 @@ def get_docai_schema_client(params, async_=True):
     return client
 
 
-def parse_input():
-    """Manage input parameters."""
-    parser = argparse.ArgumentParser(description="", add_help=False)
-    parser.add_argument(
-        "--scope",
-        type=str,
-        dest="scope",
-        required=False,
-        help="Whether the function should 'upload' or 'download' documents",
-    )
-    parser.add_argument(
-        "--document_name",
-        type=str,
-        dest="document_name",
-        required=False,
-        help="Category of the document (e.g., 'commercialInvoice', 'packingList')",
-    )
-    parser.add_argument(
-        "--for_combinations",
-        type=bool,
-        default=False,
-        dest="for_combinations",
-        required=False,
-        help="A flag to download documents into a special subfolder",
-    )
-    parser.add_argument(
-        "--n_samples",
-        type=int,
-        default=50,
-        dest="n_samples",
-        required=False,
-        help="A number of samples to download",
-    )
-
-    # Remove declared missing arguments (e.g. model_type)
-    args = vars(parser.parse_args())
-    args_no_null = {
-        k: v.split(",") if isinstance(v, str) else v
-        for k, v in args.items()
-        if v is not None
-    }
-    return args_no_null
-
-
 def setup_params(args=None):
     """
     Set up the application parameters.

src/utils.py

@@ -10,6 +10,7 @@ from typing import Literal
 
 import openpyxl
 import pandas as pd
+import requests
 from google.cloud import documentai_v1beta3 as docu_ai_beta
 from PyPDF2 import PdfReader, PdfWriter
 
@@ -364,3 +365,50 @@ def extract_top_pages(pdf_bytes, num_pages=4):
     writer.write(output)
 
     return output.getvalue()
+
+
+def get_tms_mappings(
+    input_list: list[str], embedding_type: str, llm_ports: list[str] = None
+):
+    """Get TMS mappings for the given values.
+
+    Args:
+        input_list (list[str]): List of strings to get embeddings for.
+        embedding_type (str): Type of embedding to use
+         (e.g., "container_types", "ports", "depots", "lineitems", "terminals").
+        llm_ports (list[str], optional): List of LLM ports to use. Defaults to None.
+
+    Returns:
+        dict: A dictionary with the mapping results.
+    """
+    # To test the API locally, port-forward the embedding service in the sandbox to 8080:80
+    # If you want to launch uvicorn from the tms-embedding repo, then use --port 8080 in the config file
+    base_url = (
+        "http://0.0.0.0:8080/"
+        if os.getenv("CLUSTER") is None
+        else "http://tms-embedding:80/api/v1/mappings"
+    )
+
+    # Ensure input_list is a list
+    if not isinstance(input_list, list):
+        input_list = [input_list]
+
+    # Always send a dict with named keys
+    payload = {embedding_type: input_list}
+    if llm_ports:
+        payload["llm_ports"] = llm_ports if isinstance(llm_ports, list) else [llm_ports]
+
+    # Make the POST request to the TMS mappings API
+    url = f"{base_url}/{embedding_type}"
+    response = requests.post(url=url, json=payload)
+
+    if response.status_code != 200:
+        logger.error(
+            f"Error from TMS mappings API: {response.status_code} - {response.text}"
+        )
+
+    formatted_values = (
+        response.json().get("response", {}).get("data", {}).get(input_list[0], None)
+    )
+
+    return formatted_values