handler.py

import torch
from unsloth.trainer import UnslothVisionDataCollator
from unsloth import FastVisionModel
import re
import base64
import os
from typing import Optional, Union, List, Dict, Any
if not os.path.exists('./static'): os.mkdir('./static')
import logging
from transformers import BitsAndBytesConfig


logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')


model_id = "Portx/v2-extractor-full-merged-4bit"


class PromptSet:
  main_order_information_prompt = """
      You are an expert in analyzing and extracting information from freight, shipment, or delivery orders. Please carefully read the provided order file and extract the following 11 key pieces of information. Ensure that the key names are exactly as listed below. Do not create any additional key names other than these. If any information is missing or unavailable, output '-'.
    
    #Key names and their descriptions:
    1. container_number: The container number/no of the shipment (e.g., TRKU2038448, MSDU8549321). This should be an 11-character container number, with no additional format. If not available, output '-'.
    2. bill_of_lading: The Bill of Lading number, which could include formats such as B/L No., AWS No., BL No., or ocean Bill of Lading (e.g., AXVJMER000008166, TRKU-10152009, HLCU ALY241000275). If not available, output '-'. 
    Do not extract values labeled as "HOUSE B/L NO", "House Bill of Lading", or any other variation indicating a House Bill of Lading. Only extract Master Bill of Lading numbers or those clearly marked without house references. 
    3. importing_carrier: The importing or ocean carrier, which may include SCAC codes, carrier's local agents, or sea line codes. If not available, output '-'.
    4. origin_address: The address for picking up the container, such as the origin address, pickup location, terminal, or port of discharge. Exclude loading location information. (e.g., "PORT LIBERTY NY CONTAINER TERMINAL 300 WESTERN AVE"). If not available, output '-'.
    5. destination_address: The address where the container is to be delivered, typically a company name or a specific delivery location (e.g., "AERO RECEIVING EAST, 2 BRICK PLANT ROAD, SOUTH RIVER, NJ"). If not available, output '-'.
    6. container_weight: The weight of the container (in numeric format, e.g., 58,201.44). If there are multiple weights, output the highest value. If not available, output '-'.
    7. container_weight_unit: The unit of measurement for the container's weight (e.g., LBS, KGS, KG, LB). If not available, output '-'.
    8. container_type: The type/size of the container (e.g., 40HC, 20GP FCL). If not available, output '-'.
    9. po_number: The purchase order number or customer’s PO (e.g., PO Number, customer’s PO, consol). If not available, output '-'.
    10. reference_number: The reference number, file number, or any internal reference (e.g., reference number, our ref no.). If not available, output '-'.
    11. hazmat: If the document includes any indication that the cargo is hazardous (e.g., labeled in red with "Hazardous 1", mentions of "Dangerous Goods", "Hazmat", or similar keywords), return True. If there is no indication of hazardous material or no related information is available at all, return False.
    
    #Output:
    {cnt_no: ...,
    bol: ..,
    imp_carr: ...,
    orig_addr: ...,
    dest_addr: ...,
    cnt_wt: ...,
    wt_unit: ...,
    cnt_type: ...,
    po_no: ...,
    ref_no: ...,
    hazmat:... 
    }
    
    Guidelines:
    - Very important: do not make up anything. If the information of a required field is not available, output '-' for it.
    - Output in JSON format. The JSON should contain the above 11 keys.
    """
  order_list_prompt = "How much container are there? Give to me all container numbers only in a json array?"
  multiple_container_information_prompt = "Give to me container weight, container weight unit,the container size (with type) of {query} in the same line with container_number:{query}.You must response only in a JSON format. Example output is must be 'container_number': 'OOCU6979480', 'container_type': '40HC or DV', 'weight': '46,737.52', 'weight_unit': 'LB'"



class RegexSet:
  def get_all_container_array(input_response):
    try:
      pattern = r'\[([^\]]+)\]'
      matches = re.findall(pattern, input_response)
      final_response = matches[0].split(', ')
      total_container_number = len(final_response)
      return final_response, total_container_number
    except:
      return '[]', 0


  def convert_one_order_information(input_response):
    try:
      pattern = r"'([^']+)':\s'([^']+)'"
      matches = re.findall(pattern, input_response)
      final_response = {match[0]: match[1] for match in matches}
      return final_response
    except:
      return '-'

class Utils:
    def base64_to_image(base64_string):
        image_data = base64.b64decode(base64_string)
        image = Image.open(BytesIO(image_data))
        return image

        
class EndpointHandler:
    def __init__(self, path=""):
      self.model, self.tokenizer = FastVisionModel.from_pretrained(model_id,  token=os.getenv('HF_TOKEN'),
                                                                   quantization_config=BitsAndBytesConfig(load_in_4bit=True))


    def __call__(self, data: Dict[str, bytes]) -> Dict[str, List[Any]]:
      FastVisionModel.for_inference(self.model)

      prompt, base64_image = data["inputs"]["text"], data["inputs"]["image"]

      if prompt == "0":
          final_prompt = PromptSet.main_order_information_prompt
      elif prompt == "1":
          final_prompt = PromptSet.order_list_prompt
      else:
          final_prompt = prompt

      # This already returns a PIL Image object
      image = Utils.base64_to_image(base64_image)

      messages = [
          {"role": "user", "content": [
              {"type": "image"},
              {"type": "text", "text": final_prompt}
          ]}
      ]

      input_text = self.tokenizer.apply_chat_template(messages, add_generation_prompt=False)
      inputs = self.tokenizer(image, input_text, add_special_tokens=False, return_tensors="pt").to("cuda")

      output = self.model.generate(
          **inputs,
          max_new_tokens=512,
          use_cache=True,
          temperature=1.5,
          min_p=0.9
      )

      # You should decode the output to get a human-readable result
      final_output = self.tokenizer.decode(
          output[0][len(inputs['input_ids'][0]):],
          skip_special_tokens=True
      )

      return {"predictions": final_output}