Add ONNX exports and ONNXRuntime examples

.gitattributes

@@ -3,5 +3,6 @@
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
 sp.model filter=lfs diff=lfs merge=lfs -text
 vlmo/tokenizer/sp.model filter=lfs diff=lfs merge=lfs -text

README.md

@@ -63,6 +63,75 @@ print(probs)
 `model(**inputs)` also returns `logits_per_image` and `logits_per_text`, which use the model's learned `logit_scale`.
 Those logits are useful, but they are not the same computation as the raw dot product in the original ModelScope demo.
 
+### ONNXRuntime
+
+This repo also includes two ONNX exports:
+
+- `onnx/text_encoder.onnx`
+- `onnx/image_encoder.onnx`
+
+Example:
+
+```python
+import importlib
+import json
+import os
+import sys
+
+import onnxruntime as ort
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+repo_id = "malusama/M2-Encoder-0.4B"
+model_dir = snapshot_download(repo_id=repo_id)
+sys.path.insert(0, model_dir)
+
+tokenizer_config = json.load(open(os.path.join(model_dir, "tokenizer_config.json"), "r", encoding="utf-8"))
+GLMChineseTokenizer = importlib.import_module("tokenization_glm").GLMChineseTokenizer
+M2EncoderImageProcessor = importlib.import_module("image_processing_m2_encoder").M2EncoderImageProcessor
+
+tokenizer = GLMChineseTokenizer(
+    vocab_file=os.path.join(model_dir, "sp.model"),
+    eos_token=tokenizer_config.get("eos_token"),
+    pad_token=tokenizer_config.get("pad_token"),
+    cls_token=tokenizer_config.get("cls_token"),
+    mask_token=tokenizer_config.get("mask_token"),
+    unk_token=tokenizer_config.get("unk_token"),
+)
+image_processor = M2EncoderImageProcessor.from_pretrained(model_dir)
+
+text_inputs = tokenizer(
+    ["杰尼龟", "妙蛙种子", "小火龙", "皮卡丘"],
+    padding="max_length",
+    truncation=True,
+    max_length=52,
+    return_special_tokens_mask=True,
+    return_tensors="np",
+)
+image_inputs = image_processor(Image.open("pokemon.jpeg").convert("RGB"), return_tensors="np")
+
+text_session = ort.InferenceSession(
+    os.path.join(model_dir, "onnx", "text_encoder.onnx"),
+    providers=["CPUExecutionProvider"],
+)
+image_session = ort.InferenceSession(
+    os.path.join(model_dir, "onnx", "image_encoder.onnx"),
+    providers=["CPUExecutionProvider"],
+)
+
+text_embeds = text_session.run(
+    None,
+    {
+        "input_ids": text_inputs["input_ids"],
+        "attention_mask": text_inputs["attention_mask"],
+    },
+)[0]
+image_embeds = image_session.run(
+    None,
+    {"pixel_values": image_inputs["pixel_values"]},
+)[0]
+```
+
 ## Upload
 
 Option 1:

onnx/image_encoder.onnx

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

onnx/text_encoder.onnx

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