Merge pull request #1 from sionic-ai/attention_mask

report: tf-torch mismatch

Author: sigridjineth

.gitignore

@@ -0,0 +1,3 @@
+model/
+*.zip
+__pycache__/

.python-version

@@ -0,0 +1 @@
+3.11

BGEM3WeightConverter.py

@@ -1,4 +1,5 @@
 from BGEM3TFModel import BGEM3TensorFlow, save_model_with_tokenizer
+from huggingface_hub import hf_hub_download
 
 from transformers import AutoModel
 import numpy as np
@@ -19,13 +20,33 @@ def load_sparse_weights():
 
 
 def load_colbert_weights():
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    model_path = os.path.join(script_dir, './bge-m3', 'colbert_linear.pt')
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"FileNotFoundError: {model_path}")
+    repo_id = "BAAI/bge-m3"
+    filename = "colbert_linear.pt"
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    return torch.load(model_path, map_location=device, weights_only=True)
+
+    try:
+        hf_model_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+        )
+        print(f"Loaded from Hugging Face: {hf_model_path}")
+
+        model = torch.load(hf_model_path, map_location=device)
+
+    except Exception as e:
+        print("Failed to load from Hugging Face. Reason:", e)
+
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        local_model_path = os.path.join(script_dir, 'bge-m3', 'colbert_linear.pt')
+
+        if not os.path.exists(local_model_path):
+            raise FileNotFoundError(f"FileNotFoundError: {local_model_path}")
+
+        print(f"Loaded from local: {local_model_path}")
+        model = torch.load(local_model_path, map_location=device)
+
+    return model
 
 
 def _init_colbert_weights(tf_model):
@@ -225,7 +246,7 @@ def convert_and_save_model(model_name: str, save_path: str):
 
 
 if __name__ == "__main__":
-    model_name = "./bge-m3"
+    model_name = "BAAI/bge-m3"
     save_path = "./converted_bge_m3"
 
     tf_model = convert_and_save_model(model_name, save_path)

README.md

@@ -43,21 +43,22 @@ git clone https://huggingface.co/BAAI/bge-m3
 
 3. Set up Python virtual environment:
 ```bash
-python3.11 -m venv venv
+uv venv
 source ./venv/bin/activate
-pip install -r requirements.txt
+uv sync
+# uv add -r requirements.txt
 ```
 
 ## Usage
 
 1. Convert the model weights:
 ```bash
-python BGEM3WeightConverter.py
+uv run BGEM3WeightConverter.py
 ```
 
 2. Validate the converted model:
 ```bash
-python model_conversion_validator.py
+uv run model_conversion_validator.py
 ```
 
 ## Model Architecture

model_conversion_validator.py

@@ -261,11 +261,11 @@ def compare_layer_outputs1(pt_all_layer_outputs, tf_all_layer_outputs):
 
 def main():
     # 경로 설정 (예: ./bge-m3, ./converted_bge_m3)
-    model_name_or_path = "./bge-m3"  # PyTorch 원본
+    model_name_or_path = "BAAI/bge-m3"  # PyTorch 원본
     saved_model_dir = "./converted_bge_m3"  # TF 변환본
 
     queries = [
-        "이 모델은 무엇을 하는 모델인가요?",
+        "이 모델은 무엇을 하는 모델인가요?이 모델은 무엇을 하는 모델인가요?이 모델은 무엇을 하는 모델인가요?이 모델은 무엇을 하는 모델인가요?이 모델은 무엇을 하는 모델인가요?이 모델은 무엇을 하는 모델인가요?",
         "이 모델은 무엇을 하는 모델인가요?"
     ]
 

pyproject.toml

@@ -0,0 +1,64 @@
+[project]
+name = "bge-m3-model-converter"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "absl-py==2.1.0",
+    "astunparse==1.6.3",
+    "certifi==2024.12.14",
+    "charset-normalizer==3.4.1",
+    "filelock==3.16.1",
+    "flatbuffers==24.12.23",
+    "fsspec==2024.12.0",
+    "gast==0.6.0",
+    "gdown>=5.2.0",
+    "google-pasta==0.2.0",
+    "grpcio==1.68.1",
+    "h5py==3.12.1",
+    "huggingface-hub==0.27.0",
+    "idna==3.10",
+    "jinja2==3.1.5",
+    "keras==3.7.0",
+    "libclang==18.1.1",
+    "loguru>=0.7.3",
+    "markdown==3.7",
+    "markdown-it-py==3.0.0",
+    "markupsafe==3.0.2",
+    "mdurl==0.1.2",
+    "ml-dtypes==0.3.2",
+    "mpmath==1.3.0",
+    "namex==0.0.8",
+    "networkx==3.4.2",
+    "numpy==1.26.4",
+    "nvitop>=1.4.0",
+    "opt-einsum==3.4.0",
+    "optree==0.13.1",
+    "packaging==24.2",
+    "protobuf==4.25.5",
+    "pygments==2.18.0",
+    "pyyaml==6.0.2",
+    "regex==2024.11.6",
+    "requests==2.32.3",
+    "rich==13.9.4",
+    "safetensors==0.4.5",
+    "six==1.17.0",
+    "sympy==1.13.1",
+    "tensorboard==2.16.2",
+    "tensorboard-data-server==0.7.2",
+    "tensorflow==2.16.2",
+    "tensorflow-io-gcs-filesystem==0.37.1",
+    "termcolor==2.5.0",
+    "tf-keras==2.16.0",
+    "tokenizers==0.21.0",
+    "torch==2.5.1",
+    "torchaudio>=2.5.1",
+    "torchvision>=0.20.1",
+    "tqdm==4.67.1",
+    "transformers==4.47.1",
+    "typing-extensions==4.12.2",
+    "urllib3==2.3.0",
+    "werkzeug==3.1.3",
+    "wrapt==1.17.0",
+]

torch_tf_validator.py

@@ -0,0 +1,94 @@
+import loguru 
+
+from transformers import AutoTokenizer, AutoModel
+import tensorflow as tf 
+import torch 
+
+def load_torch_model(model_path):
+    model = AutoModel.from_pretrained(model_path)
+    return model
+
+
+def load_tf_model(model_path):
+    with tf.device("/CPU:0"):
+        model = tf.saved_model.load(model_path)
+    return model
+
+
+def load_tokenizer(model_path):
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    return tokenizer
+
+
+def tokenize_wo_padding(tokenizer, text, return_tensors="pt"):
+    return tokenizer(text, padding=False, return_tensors=return_tensors)
+
+
+def tokenize_w_padding(tokenizer, text, return_tensors="pt", max_length=512):
+    return tokenizer(text, padding="max_length", max_length=max_length, return_tensors=return_tensors)
+
+
+def main():
+    # Load the model
+    model_path = "BAAI/bge-m3"
+    model_path_tf = "/workspace/BGE-M3-Model-Converter/model"
+    model = load_torch_model(model_path)
+    tokenizer = load_tokenizer(model_path)
+
+    # Tokenize the text
+    text = "Hello, my dog is cute"
+    inputs = tokenize_wo_padding(tokenizer, text)
+    inputs_w_padding = tokenize_w_padding(tokenizer, text)
+
+    # Get the output from the model
+    loguru.logger.info("Torch] Model output".ljust(50, "-"))
+    model.eval().to("cuda")
+    with torch.no_grad():
+        inputs = {k: v.to("cuda") for k, v in inputs.items()}
+        inputs_w_padding = {k: v.to("cuda") for k, v in inputs_w_padding.items()}
+
+        output = model(**inputs)
+        output_w_padding = model(**inputs_w_padding)
+        loguru.logger.info("output without padding (GT)".ljust(50, "-"))
+        loguru.logger.info(output['last_hidden_state'][:, 0])
+        loguru.logger.info("="*50)
+        loguru.logger.info("output with padding".ljust(50, "-"))
+        loguru.logger.info(output_w_padding['last_hidden_state'][:, 0])
+        loguru.logger.info("="*50)
+        err = torch.abs(output['last_hidden_state'][:, 0] - output_w_padding['last_hidden_state'][:, 0])
+        loguru.logger.info("Error".ljust(50, "-"))
+        loguru.logger.info(err.mean())
+    
+    inputs_tf = tokenize_wo_padding(tokenizer, text, return_tensors="tf")
+    inputs_tf_w_padding = tokenize_w_padding(tokenizer, text, return_tensors="tf")
+    inputs_tf_w_padding_attnFixed = inputs_tf_w_padding.copy()
+    inputs_tf_w_padding_attnFixed['attention_mask'] = tf.where(inputs_tf_w_padding['attention_mask'] == 0, -9999999, 0)
+    tf_model = load_tf_model(model_path_tf).signatures["serving_default"]
+
+    loguru.logger.info("Tensorflow] Model output".ljust(50, "-"))
+    with tf.device("/GPU:0"):
+        output_tf = tf_model(**inputs_tf)
+        output_tf_w_padding = tf_model(**inputs_tf_w_padding)
+        output_tf_w_padding_attnFixed = tf_model(**inputs_tf_w_padding_attnFixed)
+        loguru.logger.info("output without padding (GT)".ljust(50, "-"))
+        loguru.logger.info(output_tf['hidden_states'][-1][:,0])
+        loguru.logger.info("="*50)
+        loguru.logger.info("output with padding".ljust(50, "-"))
+        loguru.logger.info(output_tf_w_padding['hidden_states'][-1][:,0])
+        loguru.logger.info("="*50)
+        loguru.logger.info("output with padding (attention fixed)".ljust(50, "-"))
+        loguru.logger.info(output_tf_w_padding_attnFixed['hidden_states'][-1][:,0])
+        loguru.logger.info("="*50)
+        err_tf = tf.abs(output_tf['hidden_states'][-1][:,0] - output_tf_w_padding['hidden_states'][-1][:,0])
+        loguru.logger.info("Error".ljust(50, "-"))
+        loguru.logger.info(tf.reduce_mean(err_tf))
+        loguru.logger.info("="*50)
+        err_tf_attnFixed = tf.abs(output_tf_w_padding['hidden_states'][-1][:,0] - output_tf_w_padding_attnFixed['hidden_states'][-1][:,0])
+        loguru.logger.info("Error (attention fixed)".ljust(50, "-"))
+        loguru.logger.info(tf.reduce_mean(err_tf_attnFixed))
+        loguru.logger.info("="*50)
+        
+
+
+if __name__ == "__main__":
+    main()