handler.py

import sys
from subprocess import run


#run("pip install unsloth", shell=True, check=True)
run("pip uninstall unsloth unsloth_zoo -y", shell=True, check=True)
run("pip install unsloth unsloth_zoo --no-cache-dir --upgrade", shell=True, check=True)
#run("pip uninstall unsloth -y && pip install --upgrade --no-cache-dir --no-deps git+https://github.com/unslothai/unsloth.git", shell=True, check=True)
run("pip install uvicorn fastapi python-multipart", shell=True, check=True)
run("pip install PyMuPDF pypdf", shell=True, check=True)

import torch
from unsloth.trainer import UnslothVisionDataCollator
from unsloth import FastVisionModel
from PIL import Image
import re
import json
from fastapi import FastAPI, HTTPException, Query, Request, File, UploadFile
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import shutil
import os
import pymupdf
from pypdf import PdfReader
from enum import Enum
from pydantic import BaseModel, Field
from typing import Optional, Union, List, Dict, Any
if not os.path.exists('./static'): os.mkdir('./static')
import logging
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig


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

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_id = "Portx/v2-trial-3"


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_jpg(base64_string):
        image_data = base64.b64decode(base64_string)
        with open("./do_img.jpg", 'wb') as f:
            f.write(image_data)





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))


    model, tokenizer = FastVisionModel.from_pretrained(
                model_id,
                load_in_4bit = True, 
                use_gradient_checkpointing = "unsloth", 
                token=os.getenv('HF_TOKEN')
            )

    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

        converted_image = Utils.base64_to_jpg(base64_image)
        image = Image.open(converted_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)
        #final_output = self.tokenizer.decode(output[0][len(inputs['input_ids'][0]):], skip_special_tokens=True)
        #response = RegexSet.convert_one_order_information(input_response=final_output)

        return {"predictions": output}