hjxdl 0.2.29__py3-none-any.whl → 0.2.30__py3-none-any.whl

hdl/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.29'
-__version_tuple__ = version_tuple = (0, 2, 29)
+__version__ = version = '0.2.30'
+__version_tuple__ = version_tuple = (0, 2, 30)

hdl/utils/llm/vis.py

@@ -10,8 +10,10 @@ import torch
 import numpy as np
 from PIL import Image
 # from transformers import ChineseCLIPProcessor, ChineseCLIPModel
+from transformers import AutoModel
+from transformers import AutoTokenizer
 import open_clip
-import natsort
+# import natsort
 from redis.commands.search.field import VectorField
 from redis.commands.search.indexDefinition import IndexDefinition, IndexType
 from hdl.jupyfuncs.show.pbar import tqdm
@@ -96,7 +98,7 @@ def imgbase64_to_pilimg(img_base64: str):
         BytesIO(
             base64.b64decode(img_base64.split(",")[-1])
         )
-    )
+    ).convert('RGB')
     return img_pil
 
 
@@ -124,6 +126,7 @@ class ImgHandler:
         model_path,
         conn=None,
         model_name: str = None,
+        model_type: str = "openclip",
         device: str = "cpu",
         num_vec_dim: int = None,
         load_model: bool = True,
@@ -147,6 +150,7 @@ class ImgHandler:
         self.device = torch.device(device)
         self.model_path = model_path
         self.model_name = model_name
+        self.model_type = model_type
 
         self.db_conn = conn
         self.num_vec_dim = num_vec_dim
@@ -164,32 +168,46 @@ class ImgHandler:
         Returns:
             None
         """
-        ckpt_file = (
-            Path(self.model_path) / Path("open_clip_pytorch_model.bin")
-        ).as_posix()
-        self.open_clip_cfg = json.load(
-            open(Path(self.model_path) / Path("open_clip_config.json"))
-        )
 
-        if self.model_name is None:
-            self.model_name = (
-                self.open_clip_cfg['model_cfg']['text_cfg']['hf_tokenizer_name']
-                .split('/')[-1]
+        if self.model_type == "cpm":
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_path,
+                trust_remote_code=True
+            )
+            self.model = AutoModel.from_pretrained(
+                self.model_path,
+                trust_remote_code=True
             )
+            self.model.to(self.device)
+            self.num_vec_dim = 2304
+
+        elif self.model_type == "openclip":
+            ckpt_file = (
+                Path(self.model_path) / Path("open_clip_pytorch_model.bin")
+            ).as_posix()
+            self.open_clip_cfg = json.load(
+                open(Path(self.model_path) / Path("open_clip_config.json"))
+            )
+
+            if self.model_name is None:
+                self.model_name = (
+                    self.open_clip_cfg['model_cfg']['text_cfg']['hf_tokenizer_name']
+                    .split('/')[-1]
+                )
 
-        self.model, self.preprocess_train, self.preprocess_val = (
-            open_clip.create_model_and_transforms(
-                model_name=self.model_name,
-                pretrained=ckpt_file,
-                device=self.device,
-                # precision=precision
+            self.model, self.preprocess_train, self.preprocess_val = (
+                open_clip.create_model_and_transforms(
+                    model_name=self.model_name,
+                    pretrained=ckpt_file,
+                    device=self.device,
+                    # precision=precision
+                )
+            )
+            if self.num_vec_dim is None:
+                self.num_vec_dim = self.open_clip_cfg["model_cfg"]["embed_dim"]
+            self.tokenizer = open_clip.get_tokenizer(
+                HF_HUB_PREFIX + self.model_path
             )
-        )
-        if self.num_vec_dim is None:
-            self.num_vec_dim = self.open_clip_cfg["model_cfg"]["embed_dim"]
-        self.tokenizer = open_clip.get_tokenizer(
-            HF_HUB_PREFIX + self.model_path
-        )
 
     def get_img_features(
         self,
@@ -222,16 +240,24 @@ class ImgHandler:
                     f"Not supported image type for {type(img)}"
                 )
 
-
-        with torch.no_grad(), torch.amp.autocast(self.device_str):
-            imgs = torch.stack([
-                self.preprocess_val(image).to(self.device)
-                for image in images_fixed
-            ])
-            img_features = self.model.encode_image(imgs, **kwargs)
-            img_features /= img_features.norm(dim=-1, keepdim=True)
-            if to_numpy:
-                img_features = img_features.cpu().numpy()
+        if self.model_type == "cpm":
+            with torch.no_grad():
+                img_features = self.model(
+                    text=[""] * len(images_fixed),
+                    image=images_fixed,
+                    tokenizer=self.tokenizer
+                ).reps
+
+        if self.model_type == "openclip":
+            with torch.no_grad(), torch.amp.autocast(self.device_str):
+                imgs = torch.stack([
+                    self.preprocess_val(image).to(self.device)
+                    for image in images_fixed
+                ])
+                img_features = self.model.encode_image(imgs, **kwargs)
+                img_features /= img_features.norm(dim=-1, keepdim=True)
+        if to_numpy:
+            img_features = img_features.cpu().numpy()
         return img_features
 
     def get_text_features(
@@ -253,17 +279,25 @@ class ImgHandler:
         Example:
             get_text_features(["text1", "text2"], to_numpy=True)
         """
-        with torch.no_grad(), torch.amp.autocast(self.device_str):
-            txts = self.tokenizer(
-                texts,
-                context_length=self.model.context_length
-            ).to(self.device)
-            txt_features = self.model.encode_text(txts, **kwargs)
-            txt_features /= txt_features.norm(dim=-1, keepdim=True)
-            if to_numpy:
-                txt_features = txt_features.cpu().numpy()
-        return txt_features
 
+        if self.model_type == "cpm":
+            with torch.no_grad():
+                txt_features = self.model(
+                    text=texts,
+                    image=[None] * len(texts),
+                    tokenizer=self.tokenizer
+                ).reps
+        elif self.model_type == "openclip":
+            with torch.no_grad(), torch.amp.autocast(self.device_str):
+                txts = self.tokenizer(
+                    texts,
+                    context_length=self.model.context_length
+                ).to(self.device)
+                txt_features = self.model.encode_text(txts, **kwargs)
+                txt_features /= txt_features.norm(dim=-1, keepdim=True)
+        if to_numpy:
+            txt_features = txt_features.cpu().numpy()
+        return txt_features
 
     def get_text_img_probs(
         self,

{hjxdl-0.2.29.dist-info → hjxdl-0.2.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hjxdl
-Version: 0.2.29
+Version: 0.2.30
 Summary: A collection of functions for Jupyter notebooks
 Home-page: https://github.com/huluxiaohuowa/hdl
 Author: Jianxing Hu

{hjxdl-0.2.29.dist-info → hjxdl-0.2.30.dist-info}/RECORD

@@ -1,5 +1,5 @@
 hdl/__init__.py,sha256=GffnD0jLJdhkd-vo989v40N90sQbofkayRBwxc6TVhQ,72
-hdl/_version.py,sha256=AtwvoTC96AXOg97Emp1_Wmo7L-xUzOR7aFTUto-9JfA,413
+hdl/_version.py,sha256=YLA3hy-44LcRh6Hcq_A-SyVyfXwkdObMoIBSVcngmDI,413
 hdl/args/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/args/loss_args.py,sha256=s7YzSdd7IjD24rZvvOrxLLFqMZQb9YylxKeyelSdrTk,70
 hdl/controllers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -132,13 +132,13 @@ hdl/utils/llm/chatgr.py,sha256=GO2G7g6YybduA5VCUuGjvEsJfC_6L7rycSnPeHMcxyM,2820
 hdl/utils/llm/embs.py,sha256=Tf0FOYrOFZp7qQpEPiSCXzlgyHH0X9HVTUtsup74a9E,7174
 hdl/utils/llm/extract.py,sha256=2sK_WJzmYIc8iuWaM9DA6Nw3_6q1O4lJ5pKpcZo-bBA,6512
 hdl/utils/llm/llama_chat.py,sha256=watcHGOaz-bv3x-yDucYlGk5f8FiqfFhwWogrl334fk,4387
-hdl/utils/llm/vis.py,sha256=Kixrhc3eByHjdxiAcB2BnthsS31dHMavp6qc2JX46Dc,16292
+hdl/utils/llm/vis.py,sha256=RWeI6lSmzCDG2HJMq8-teuC7to4pPiR0ee2Hx1clbRw,17656
 hdl/utils/llm/visrag.py,sha256=_PuKtmQIXD5bnmXwDWhTLdzOhgC42JiqdMNb1uKA7n8,9190
 hdl/utils/schedulers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/schedulers/norm_lr.py,sha256=bDwCmdEK-WkgxQMFBiMuchv8Mm7C0-GZJ6usm-PQk14,4461
 hdl/utils/weather/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/weather/weather.py,sha256=k11o6wM15kF8b9NMlEfrg68ak-SfSYLN3nOOflFUv-I,4381
-hjxdl-0.2.29.dist-info/METADATA,sha256=7qiDTPY06pouX5FgbV4ekWvdDdXKae-Q9m1iGrIO1GA,836
-hjxdl-0.2.29.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-hjxdl-0.2.29.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
-hjxdl-0.2.29.dist-info/RECORD,,
+hjxdl-0.2.30.dist-info/METADATA,sha256=aKoW2m7jTjCWQSTm4Kzxi7n5YWyXXxjHlmqHjZjG8D4,836
+hjxdl-0.2.30.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+hjxdl-0.2.30.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
+hjxdl-0.2.30.dist-info/RECORD,,

{hjxdl-0.2.29.dist-info → hjxdl-0.2.30.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.2.0)
+Generator: setuptools (75.3.0)
 Root-Is-Purelib: true
 Tag: py3-none-any