Update handler.py

handler.py

@@ -1,48 +1,134 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
- 
+from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
+import sys
+from subprocess import run
+from PIL import Image
+
+run("pip install flash-attn --no-build-isolation", shell=True, check=True)
+model_id = "ibm-granite/granite-vision-3.1-2b-preview"
+
+bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        llm_int8_skip_modules=["vision_tower", "lm_head"],
+        llm_int8_enable_fp32_cpu_offload=True
+)
+
+
+try:
+    import flash_attn
+    print("FlashAttention is installed")
+    USE_FLASH_ATTENTION = True
+except ImportError:
+    print("FlashAttention is not installed")
+    USE_FLASH_ATTENTION = False
+
+
 # check for GPU
 device = 0 if torch.cuda.is_available() else -1
- 
-# multi-model list
-multi_model_list = [
-    {"model_id": "distilbert-base-uncased-finetuned-sst-2-english", "task": "text-classification"},
-    {"model_id": "Helsinki-NLP/opus-mt-en-de", "task": "translation"},
-    {"model_id": "facebook/bart-large-cnn", "task": "summarization"},
-    {"model_id": "dslim/bert-base-NER", "task": "token-classification"},
-    {"model_id": "textattack/bert-base-uncased-ag-news", "task": "text-classification"},
-    {"model_id": "openai-community/gpt2","task": "text-generation"}
-]
- 
+
+
+
+inputs = processor.apply_chat_template(
+    conversation,
+    add_generation_prompt=True,
+    tokenize=True,
+    return_dict=True,
+    return_tensors="pt"
+).to(device)
+
+
+class Utils:
+    def convert_base64_to_jpg(base64_string):
+        image_data = base64.b64decode(base64_string)
+        with open("./do_img.jpg", 'wb') as f:
+            f.write(image_data)
+
+
+class PromptSet:
+    system_message = "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 10 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 '-'."
+  main_order_information_prompt = """
+  #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 '-'.
+    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 '-'.
+  #Output:
+    {container_number: ...,
+    bill_of_lading: ..,
+    importing_carrier: ...,
+    origin_address: ...,
+    destination_address: ...,
+    container_weight: ...,
+    container_weight_unit: ...,
+    container_type: ...,
+    po_number: ...,
+    reference_number: ...
+    }
+  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 10 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 EndpointHandler():
     def __init__(self, path=""):
-        self.multi_model={}
-        # load all the models onto device
-        for model in multi_model_list:
-            self.multi_model[model["model_id"]] = pipeline(model["task"], model=model["model_id"], device=device)
- 
+        self.model=AutoModelForVision2Seq.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16,
+                                                          quantization_config=bnb_config,
+                                                          _attn_implementation="flash_attention_2" if USE_FLASH_ATTENTION else None)
+        self.processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
+
     def __call__(self, data):
         # deserialize incomin request
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
-        model_id = data.pop("model_id", None)
         prompt_id = data.pop("prompt_id", None)
+        base64_image = data.pop("image", None)
+
+        converted_image = Utils.convert_base64_to_jpg(base64_image)
 
+        
         if prompt_id==1:
-            inputs="What is the ai in 3 sentences?"
+            final_prompt=PromptSet.main_order_information_prompt
         elif prompt_id==2:
-            inputs="Who is elon musk?"
+            final_prompt=PromptSet.order_list_prompt
+        elif prompt_id==3:
+            final_prompt=PromptSet.multiple_container_information_prompt
         else:
             pass
             
-        # check if model_id is in the list of models
-        if model_id is None or model_id not in self.multi_model:
-            raise ValueError(f"model_id: {model_id} is not valid. Available models are: {list(self.multi_model.keys())}")
- 
-        # pass inputs with all kwargs in data
-        if parameters is not None:
-            prediction = self.multi_model[model_id](inputs, **parameters)
-        else:
-            prediction = self.multi_model[model_id](inputs)
-        # postprocess the prediction
+
+
+        conversation = [{
+            "role": "system",
+            "content": [
+                {
+                    "type": "text",
+                    "text": PromptSet.system_message
+                }
+            ],
+        },{
+        "role": "user",
+        "content": [
+            {"type": "image", "url": "./do_img.jpg"},
+            {"type": "text", "text": final_prompt},
+        ],},
+                       ]
+
+        model_inputs = self.processor.apply_chat_template(conversation,add_generation_prompt=True,
+                                                          tokenize=True, return_dict=True,return_tensors="pt").to(device)
+        
+
+        output = model.generate(**model_inputs, max_new_tokens=512)
+        prediction = processor.decode(output[0], skip_special_tokens=True)
         return prediction