various type fixes

Author: fabnemEPFL

src/mmore/colpali/retriever.py

@@ -62,9 +62,11 @@ def load_model(model_name: str, device: str):
         Tuple of (model, processor) ready for inference
     """
     logger.info(f"Loading ColPali model: {model_name}")
+
+    bfloat16: torch.dtype = torch.bfloat16
     model = ColPali.from_pretrained(
         model_name,
-        torch_dtype=torch.bfloat16,
+        torch_dtype=bfloat16,
         device_map=device,
     ).eval()
     processor = ColPaliProcessor.from_pretrained(model_name)
@@ -95,7 +97,7 @@ def embed_queries(texts: List[str], model, processor) -> List[np.ndarray]:
         with torch.no_grad():
             batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
             emb = model(**batch_query)
-            vectors.extend(list(torch.unbind(emb.to("cpu"))))
+            vectors.extend(list(emb.to("cpu").unbind()))
     return [v.float().numpy() for v in vectors]
 
 

src/mmore/colpali/run_process.py

@@ -86,9 +86,9 @@ def cleanup(self):
 class ColPaliEmbedder:
     def __init__(self, model_name: str = "vidore/colpali-v1.3", device: str = "cuda:0"):
         self.device = device
-        dtype = torch.bfloat16
+        bfloat16: torch.dtype = torch.bfloat16
         self.model = ColPali.from_pretrained(
-            model_name, torch_dtype=dtype, device_map=device
+            model_name, torch_dtype=bfloat16, device_map=device
         ).eval()
         self.processor = ColPaliProcessor.from_pretrained(model_name)
 
@@ -112,7 +112,7 @@ def embed_images(
             with torch.no_grad():
                 batch_doc = {k: v.to(self.model.device) for k, v in batch_doc.items()}
                 embeddings_doc = self.model(**batch_doc)
-            ds.extend(list(torch.unbind(embeddings_doc.to(self.device))))
+            ds.extend(list(embeddings_doc.to(self.device).unbind()))
         ds_np = [d.float().cpu().numpy() for d in ds]
         return ds_np
 

src/mmore/process/processors/base.py

@@ -216,12 +216,13 @@ def __getstate__(self):
             del state["_pool"]
         return state
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: Dict[str, Any]):
         """
         Called when the object is unpickled (received by the worker).
         We restore the state and set _pool to None (workers don't need the pool manager).
         """
-        self.__dict__.update(state)
+        for key, value in state.items():
+            setattr(self, key, value)
         # Initialize _pool as None in the worker process
         self._pool = None
         # Workers should never own the pool

src/mmore/process/processors/media_processor.py

@@ -1,12 +1,14 @@
 import logging
 import os
 import tempfile
-from typing import List
+from typing import List, cast
 
+import numpy as np
 import torch
 from moviepy.audio.io.AudioFileClip import AudioFileClip
 from moviepy.video.io.VideoFileClip import VideoFileClip
 from PIL import Image
+from torch._C import device as torch_device
 from transformers.pipelines import pipeline as pipeline_t
 
 from ...type import FileDescriptor, MultimodalSample
@@ -19,10 +21,10 @@ class MediaProcessor(Processor):
     @staticmethod
     def _get_available_devices():
         if torch.cuda.is_available():
-            return [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
+            return [torch_device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
         if torch.backends.mps.is_available():
-            return [torch.device("mps")]
-        return [torch.device("cpu")]
+            return [torch_device("mps")]
+        return [torch_device("cpu")]
 
     devices = _get_available_devices()
     pipelines = []
@@ -155,7 +157,7 @@ def _extract_video_frames(file_path: str) -> List[Image.Image]:
                     for i in range(num_thumbnails):
                         t = min(i * sample_rate, duration - 0.1)
                         frame = clip.get_frame(t)
-                        image = Image.fromarray(frame).convert("RGB")
+                        image = Image.fromarray(cast(np.ndarray, frame).convert("RGB"))
                         images.append(image)
                 logger.info(f"Extracted {len(images)} images from {file_path}.")
             except Exception as e:

src/mmore/process/processors/pdf_processor.py

@@ -2,7 +2,7 @@
 import logging
 import re
 from multiprocessing import Manager, Process, set_start_method
-from typing import List, Optional, Tuple, cast
+from typing import Any, Dict, List, Optional, Tuple, cast
 
 import pymupdf
 import torch
@@ -148,7 +148,7 @@ def process(self, file_path: str) -> MultimodalSample:
 
         paragraph_starts, text = self._parse_pagination(cast(str, text))
 
-        metadata = {"file_path": file_path}
+        metadata: Dict[str, Any] = {"file_path": file_path}
         if paragraph_starts:
             metadata["paragraph_starts"] = paragraph_starts
 
@@ -218,7 +218,7 @@ def _extract_images(pdf_doc, xref) -> Optional[Image.Image]:
                 if image_bytes is None:
                     logging.error(f"No image data found for xref {xref}")
 
-                return Image.open(io.BytesIO(image_bytes)).convert("RGB")
+                return Image.open(io.BytesIO(cast(bytes, image_bytes))).convert("RGB")
 
             except KeyError as e:
                 logging.error(f"KeyError while extracting image: {e}")
@@ -236,7 +236,7 @@ def _extract_images(pdf_doc, xref) -> Optional[Image.Image]:
                 )
                 return None
 
-        for page_num, page in enumerate(pdf_doc):
+        for page_num, page in enumerate(pdf_doc):  # pyright: ignore[reportArgumentType]
             text = clean_text(page.get_text())  # type: ignore[attr-defined]
 
             if text.strip():

src/mmore/rag/model/dense/multimodal.py

@@ -14,8 +14,9 @@
 class MultimodalEmbeddings(Embeddings):
     def __init__(self, model_name: str):
         super().__init__()
+        float16: torch.dtype = torch.float16
         self.model = AutoModelForImageTextToText.from_pretrained(
-            model_name, torch_dtype=torch.float16, device_map="auto"
+            model_name, torch_dtype=float16, device_map="auto"
         )
         self.processor = AutoProcessor.from_pretrained(model_name)
         self.device = self.model.device