Upload 3 files

clip_for_image_classification.py

@@ -0,0 +1,56 @@
+from transformers import CLIPVisionConfig, FlaxCLIPVisionPreTrainedModel
+from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule
+import jax.numpy as jnp
+from flax import linen as nn
+import jax
+from transformers.modeling_flax_outputs import FlaxSequenceClassifierOutput
+
+
+class FlaxCLIPForImageClassificationModule(nn.Module):
+    config: CLIPVisionConfig
+    dtype: jnp.dtype = jnp.float32
+
+    def setup(self):
+        self.vit = FlaxCLIPVisionModule(config=self.config, dtype=self.dtype)
+        self.classifier = nn.Dense(
+            self.config.num_labels,
+            dtype=self.dtype,
+            kernel_init=jax.nn.initializers.variance_scaling(
+                self.config.initializer_range ** 2, "fan_in", "truncated_normal"
+            ),
+        )
+
+    def __call__(
+            self,
+            pixel_values=None,
+            deterministic: bool = True,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.vit(
+            pixel_values,
+            deterministic=deterministic,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.classifier(hidden_states[:, 0, :])
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return output
+
+        return FlaxSequenceClassifierOutput(
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class FlaxCLIPForImageClassification(FlaxCLIPVisionPreTrainedModel):
+    module_class = FlaxCLIPForImageClassificationModule

inference.py

@@ -0,0 +1,38 @@
+import json
+from clip_for_image_classification import FlaxCLIPForImageClassification
+from PIL import Image
+import jax
+import numpy as np
+from transformers import CLIPImageProcessor
+import os
+
+os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+
+model = FlaxCLIPForImageClassification.from_pretrained("Thouph/clip-vit-l-224-patch14-datacomp-image-classification")
+image_processor = CLIPImageProcessor.from_pretrained("laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K")
+image = Image.open("/your/image/here.jpg")
+inputs = image_processor(images=image, return_tensors="jax")
+outputs = model(**inputs)
+logits = outputs.logits
+probabilities = jax.nn.sigmoid(logits)
+probabilities = np.asarray(probabilities).copy()
+
+
+def topk_by_sort(input, k, axis=None, ascending=False):
+    if not ascending:
+        input *= -1
+    ind = np.argsort(input, axis=axis)
+    ind = np.take(ind, np.arange(k), axis=axis)
+    if not ascending:
+        input *= -1
+    val = np.take_along_axis(input, ind, axis=axis)
+    return ind, val
+
+
+indices, values = topk_by_sort(probabilities, 100)
+
+with open("7748tags.json", "r") as file:
+    allowed_tags = json.load(file)
+
+for index, value in zip(indices, values, strict=True):
+    print(allowed_tags[index], value)