add custom handler

.gitignore

@@ -0,0 +1,13 @@
+# gitignore
+
+.venv
+test_image_1.jpg
+test_image_2.jpg
+test_image_3.jpg
+test_image_4.jpg
+test_image_5.jpg
+test_image_6.jpg
+test_image_7.jpg
+test_image_8.jpg
+test_image_9.jpg
+

handler.py

@@ -1,33 +1,27 @@
-from PIL import Image
 import torch
+from PIL import Image
 from transformers import AutoModel, AutoTokenizer
 
-class ModelHandler:
-    def __init__(self):
-        # Load the model and tokenizer with appropriate weights
+class EndpointHandler:
+    def __init__(self, path):
         self.model = AutoModel.from_pretrained(
-            'fredaddy/MiniCPM-V-2_6', 
+            path, 
             trust_remote_code=True,
-            attn_implementation='sdpa', 
-            torch_dtype=torch.bfloat16
-        ).eval().cuda()
-        
-        self.tokenizer = AutoTokenizer.from_pretrained('fredaddy/MiniCPM-V-2_6', trust_remote_code=True)
+            attn_implementation='sdpa',
+            torch_dtype=torch.float16
+        )
+        self.model = self.model.eval().cuda()
+        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
 
-    def preprocess(self, inputs):
-        # Preprocess image input
-        image = Image.open(inputs['image'].file).convert('RGB')
-        question = inputs.get("question", "Extract all data in the image. Be extremely careful to ensure that you don't miss anything. It's imperative that you extract and digitize everything on that page.")
+    def __call__(self, data):
+        image = Image.open(data['inputs']['image'].file).convert('RGB')
+        question = data['inputs'].get("question", "Extract all data in the image. Be extremely careful to ensure that you don't miss anything. It's imperative that you extract and digitize everything on that page.")
         msgs = [{'role': 'user', 'content': [image, question]}]
-        return msgs
-
-    def inference(self, msgs):
-        # Run inference on the model
-        result = self.model.chat(image=None, msgs=msgs, tokenizer=self.tokenizer)
-        return result
 
-    def postprocess(self, result):
-        # Postprocess the output from the model
-        return {"generated_text": result}
+        res = self.model.chat(
+            image=None,
+            msgs=msgs,
+            tokenizer=self.tokenizer
+        )
 
-service = ModelHandler()
+        return {"generated_text": res}

model-00001-of-00004.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:327b5e63ac7729d00b0b17233a0d648a623fc080cc1069dce9cb79478cdaf2b5
-size 4874808328

model-00002-of-00004.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e6d7af1f12ba8c72fec16ba3f3951df5643bef24de8b0af5baec67290c07ca4f
-size 4932751496

model-00003-of-00004.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e9ba1332c74a92b9c8af5fbc2c2da1baa8cd5ed4fd8056b98a586ff4bc3e54cf
-size 4330865648

model-00004-of-00004.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:16fcfafa8caaa67e16f9dae7cc79bf081eb59d4e248723a8e79d9678e42732fe
-size 2060017080

requirements.txt

@@ -0,0 +1,6 @@
+Pillow==10.1.0
+torch==2.1.2
+torchvision==0.16.2
+transformers==4.40.0
+sentencepiece==0.1.99
+flash-attn==2.3.6

test_handler.py

@@ -0,0 +1,28 @@
+from handler import EndpointHandler
+from PIL import Image
+import io
+
+def test_endpoint():
+    # Initialize the handler
+    handler = EndpointHandler("openbmb/MiniCPM-V-2_6")
+
+    # Load a test image
+    with open("test_image.jpg", "rb") as image_file:
+        image_bytes = image_file.read()
+
+    # Create a mock request data
+    mock_data = {
+        "inputs": {
+            "image": type('MockFile', (), {'file': io.BytesIO(image_bytes)})(),
+            "question": "What is in this image?"
+        }
+    }
+
+    # Call the handler
+    result = handler(mock_data)
+
+    # Print the result
+    print(result["generated_text"])
+
+if __name__ == "__main__":
+    test_endpoint()