Docstring V2

.gitignore

@@ -1 +1,2 @@
-*.DS*
+*.DS*
+*__pycache__*

handler.py

@@ -1,63 +1,3 @@
-# import torch
-# from typing import Dict, Any
-# from PIL import Image
-# import base64
-# from io import BytesIO
-
-
-# class EndpointHandler:
-#     def __init__(self, path: str = ""):
-#         # Import your model and processor inside the class
-#         from colpali_engine.models import ColQwen2, ColQwen2Processor
-
-#         # Load the model and processor
-#         self.model = ColQwen2.from_pretrained(
-#             path,
-#             torch_dtype=torch.bfloat16,
-#         ).eval()
-#         self.processor = ColQwen2Processor.from_pretrained(path)
-
-#         # Determine the device
-#         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#         self.model.to(self.device)
-
-#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-#         # Extract images from the input data
-#         images_data = data.get("inputs", [])
-
-#         if not images_data:
-#             return {"error": "No images provided in 'inputs'."}
-
-#         # Process images
-#         images = []
-#         for img_data in images_data:
-#             if isinstance(img_data, str):
-#                 try:
-#                     # Assume base64-encoded image
-#                     image_bytes = base64.b64decode(img_data)
-#                     image = Image.open(BytesIO(image_bytes)).convert("RGB")
-#                     images.append(image)
-#                 except Exception as e:
-#                     return {"error": f"Invalid image data: {e}"}
-#             else:
-#                 return {"error": "Images should be base64-encoded strings."}
-
-#         # Prepare inputs
-#         batch_images = self.processor.process_images(images)
-
-#         # Move tensors to the device
-#         batch_images = {k: v.to(self.device) for k, v in batch_images.items()}
-
-#         # Generate embeddings
-#         with torch.no_grad():
-#             image_embeddings = self.model(**batch_images)
-
-#         # Convert embeddings to a list
-#         embeddings_list = image_embeddings.cpu().tolist()
-
-#         return {"embeddings": embeddings_list}
-
-
 import torch
 from typing import Dict, Any, List
 from PIL import Image
@@ -67,13 +7,17 @@ from io import BytesIO
 
 class EndpointHandler:
     """
-    A handler class for processing image data, generating embeddings using a specified model and processor.
+    A handler class for processing image data and generating embeddings using a specified model and processor.
 
     Attributes:
-        model: The pre-trained model used for generating embeddings.
-        processor: The pre-trained processor used to process images before model inference.
-        device: The device (CPU or CUDA) used to run model inference.
-        default_batch_size: The default batch size for processing images in batches.
+        model (:obj:):
+            The pre-trained model used for generating embeddings.
+        processor (:obj:):
+            The pre-trained processor used to process images before model inference.
+        device (:obj:):
+            The device (CPU or CUDA) used to run model inference.
+        default_batch_size (:obj:int:):
+            The default batch size for processing images in batches.
     """
 
     def __init__(self, path: str = "", default_batch_size: int = 4):
@@ -81,8 +25,13 @@ class EndpointHandler:
         Initializes the EndpointHandler with a specified model path and default batch size.
 
         Args:
-            path (str): Path to the pre-trained model and processor.
-            default_batch_size (int): Default batch size for image processing.
+            path (:obj:`str`, optional):
+                Path to the pre-trained model and processor.
+            default_batch_size (:obj:`int`, optional):
+                Default batch size for image processing.
+
+        Return:
+            None
         """
         from colpali_engine.models import ColQwen2, ColQwen2Processor
 
@@ -101,10 +50,11 @@ class EndpointHandler:
         Processes a batch of images and generates embeddings.
 
         Args:
-            images (List[Image.Image]): List of images to process.
+            images (:obj:`List[Image.Image]`):
+                List of images to process.
 
-        Returns:
-            List[List[float]]: List of embeddings for each image.
+        Return:
+            A :obj:`List[List[float]]`. A list of embeddings for each image, where each embedding is a list of floats.
         """
         batch_images = self.processor.process_images(images)
         batch_images = {k: v.to(self.device) for k, v in batch_images.items()}
@@ -119,10 +69,12 @@ class EndpointHandler:
         Processes input data containing base64-encoded images, decodes them, and generates embeddings.
 
         Args:
-            data (Dict[str, Any]): Dictionary containing input images and optional batch size.
+            data (:obj:`Dict[str, Any]`):
+                Includes the input data and the parameters for the inference, such as "inputs" containing a list of base64-encoded images and an optional "batch_size".
 
-        Returns:
-            Dict[str, Any]: Dictionary containing generated embeddings or error messages.
+        Return:
+            A :obj:`dict`. The object returned should be a dict like {"embeddings": [[0.6331314444541931, 0.8802216053009033, ..., -0.7866355180740356]]} containing:
+                - "embeddings": A list of lists, where each inner list is a set of floats corresponding to the embeddings of each image.
         """
         images_data = data.get("inputs", [])
         batch_size = data.get("batch_size", self.default_batch_size)