chore: import upstream snapshot with attribution

examples/inference/embedder/encoder_only/auto_base_multi_devices.py

@@ -0,0 +1,34 @@
+import os
+from FlagEmbedding import FlagAutoModel
+
+
+def test_base_multi_devices():
+    model = FlagAutoModel.from_finetuned(
+        'BAAI/bge-small-en-v1.5',
+        query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
+        devices=["cuda:0", "cuda:1"],   # if you don't have GPUs, you can use ["cpu", "cpu"]
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is the capital of France?",
+        "What is the population of China?",
+    ] * 100
+    passages = [
+        "Paris is the capital of France.",
+        "The population of China is over 1.4 billion people."
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(queries)
+    passages_embeddings = model.encode_corpus(passages)
+    
+    cos_scores = queries_embeddings @ passages_embeddings.T
+    print(cos_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_base_multi_devices()
+    
+    print("--------------------------------")
+    print("Expected Output:")
+    print("[[0.7944 0.4492]\n [0.5806 0.801 ]]")

examples/inference/embedder/encoder_only/auto_base_single_device.py

@@ -0,0 +1,34 @@
+import os
+from FlagEmbedding import FlagAutoModel
+
+
+def test_base_single_device():
+    model = FlagAutoModel.from_finetuned(
+        'BAAI/bge-small-en-v1.5',
+        query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
+        devices="cuda:0",   # if you don't have a GPU, you can use "cpu"
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is the capital of France?",
+        "What is the population of China?",
+    ] * 100
+    passages = [
+        "Paris is the capital of France.",
+        "The population of China is over 1.4 billion people."
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(queries)
+    passages_embeddings = model.encode_corpus(passages)
+    
+    cos_scores = queries_embeddings @ passages_embeddings.T
+    print(cos_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_base_single_device()
+    
+    print("--------------------------------")
+    print("Expected Output:")
+    print("[[0.7944 0.4492]\n [0.58   0.801 ]]")

examples/inference/embedder/encoder_only/auto_m3_multi_devices.py

@@ -0,0 +1,52 @@
+import os
+from FlagEmbedding import FlagAutoModel
+
+
+def test_m3_multi_devices():
+    model = FlagAutoModel.from_finetuned(
+        'BAAI/bge-m3',
+        devices=["cuda:0", "cuda:1"],   # if you don't have GPUs, you can use ["cpu", "cpu"]
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is BGE M3?",
+        "Defination of BM25"
+    ] * 100
+    passages = [
+        "BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 
+        "BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(
+        queries,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    passages_embeddings = model.encode_corpus(
+        passages,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    
+    dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
+    sparse_scores = model.compute_lexical_matching_score(
+        queries_embeddings["lexical_weights"],
+        passages_embeddings["lexical_weights"],
+    )
+
+    print("Dense score:\n", dense_scores[:2, :2])
+    print("Sparse score:\n", sparse_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_m3_multi_devices()
+
+    print("--------------------------------")
+    print("Expected Output:")
+    print("Dense score:")
+    print(" [[0.626  0.3477]\n [0.3499 0.678 ]]")
+    print("Sparse score:")
+    print(" [[0.19561768 0.00878906]\n [0.         0.18030453]]")

examples/inference/embedder/encoder_only/auto_m3_single_device.py

@@ -0,0 +1,52 @@
+import os
+from FlagEmbedding import FlagAutoModel
+
+
+def test_m3_single_device():
+    model = FlagAutoModel.from_finetuned(
+        'BAAI/bge-m3',
+        devices="cuda:0",   # if you don't have a GPU, you can use "cpu"
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is BGE M3?",
+        "Defination of BM25"
+    ] * 100
+    passages = [
+        "BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 
+        "BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(
+        queries,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    passages_embeddings = model.encode_corpus(
+        passages,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    
+    dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
+    sparse_scores = model.compute_lexical_matching_score(
+        queries_embeddings["lexical_weights"],
+        passages_embeddings["lexical_weights"],
+    )
+
+    print("Dense score:\n", dense_scores[:2, :2])
+    print("Sparse score:\n", sparse_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_m3_single_device()
+
+    print("--------------------------------")
+    print("Expected Output:")
+    print("Dense score:")
+    print(" [[0.626  0.3477]\n [0.3496 0.678 ]]")
+    print("Sparse score:")
+    print(" [[0.19554901 0.00880432]\n [0.         0.18036556]]")

examples/inference/embedder/encoder_only/base_multi_devices.py

@@ -0,0 +1,36 @@
+import os
+from FlagEmbedding import FlagModel
+
+
+def test_base_multi_devices():
+    model = FlagModel(
+        'BAAI/bge-small-en-v1.5',
+        query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
+        query_instruction_format="{}{}",
+        devices=["cuda:0", "cuda:1"],   # if you don't have GPUs, you can use ["cpu", "cpu"]
+        pooling_method='cls',
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is the capital of France?",
+        "What is the population of China?",
+    ] * 100
+    passages = [
+        "Paris is the capital of France.",
+        "The population of China is over 1.4 billion people."
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(queries)
+    passages_embeddings = model.encode_corpus(passages)
+    
+    cos_scores = queries_embeddings @ passages_embeddings.T
+    print(cos_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_base_multi_devices()
+    
+    print("--------------------------------")
+    print("Expected Output:")
+    print("[[0.7944 0.4492]\n [0.5806 0.801 ]]")

examples/inference/embedder/encoder_only/base_single_device.py

@@ -0,0 +1,36 @@
+import os
+from FlagEmbedding import FlagModel
+
+
+def test_base_single_device():
+    model = FlagModel(
+        'BAAI/bge-small-en-v1.5',
+        query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
+        query_instruction_format="{}{}",
+        devices="cuda:0",   # if you don't have a GPU, you can use "cpu"
+        pooling_method='cls',
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is the capital of France?",
+        "What is the population of China?",
+    ] * 100
+    passages = [
+        "Paris is the capital of France.",
+        "The population of China is over 1.4 billion people."
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(queries)
+    passages_embeddings = model.encode_corpus(passages)
+    
+    cos_scores = queries_embeddings @ passages_embeddings.T
+    print(cos_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_base_single_device()
+    
+    print("--------------------------------")
+    print("Expected Output:")
+    print("[[0.7944 0.4492]\n [0.58   0.801 ]]")

examples/inference/embedder/encoder_only/m3_multi_devices.py

@@ -0,0 +1,53 @@
+import os
+from FlagEmbedding import BGEM3FlagModel
+
+
+def test_m3_multi_devices():
+    model = BGEM3FlagModel(
+        'BAAI/bge-m3',
+        devices=["cuda:0", "cuda:1"],   # if you don't have GPUs, you can use ["cpu", "cpu"]
+        pooling_method='cls',
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is BGE M3?",
+        "Defination of BM25"
+    ] * 100
+    passages = [
+        "BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 
+        "BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(
+        queries,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    passages_embeddings = model.encode_corpus(
+        passages,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    
+    dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
+    sparse_scores = model.compute_lexical_matching_score(
+        queries_embeddings["lexical_weights"],
+        passages_embeddings["lexical_weights"],
+    )
+
+    print("Dense score:\n", dense_scores[:2, :2])
+    print("Sparse score:\n", sparse_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_m3_multi_devices()
+
+    print("--------------------------------")
+    print("Expected Output:")
+    print("Dense score:")
+    print(" [[0.626  0.3477]\n [0.3499 0.678 ]]")
+    print("Sparse score:")
+    print(" [[0.19561768 0.00878906]\n [0.         0.18030453]]")

examples/inference/embedder/encoder_only/m3_multi_devices_compute_score.py

@@ -0,0 +1,55 @@
+import os
+from FlagEmbedding import BGEM3FlagModel
+
+
+def test_m3_multi_devices():
+    model = BGEM3FlagModel(
+        'BAAI/bge-m3',
+        devices=["cuda:0", "cuda:1"],   # if you don't have GPUs, you can use ["cpu", "cpu"]
+        pooling_method='cls',
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is BGE M3?",
+        "Defination of BM25"
+    ] * 100
+    passages = [
+        "BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 
+        "BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
+    ] * 100
+    
+    sentence_pairs = list(zip(queries, passages))
+    scores_dict = model.compute_score(
+        sentence_pairs,
+        weights_for_different_modes=[1., 0.3, 1.]
+    )
+    
+    queries.reverse()
+    sentence_pairs = list(zip(queries, passages))
+    
+    scores_dict_reverse = model.compute_score(
+        sentence_pairs,
+        weights_for_different_modes=[1., 0.3, 1.]
+    )
+    
+    scores_dict = {
+        key: value[:2]
+        for key, value in scores_dict.items()
+    }
+    scores_dict_reverse = {
+        key: value[:2]
+        for key, value in scores_dict_reverse.items()
+    }
+    
+    print(scores_dict)
+    print(scores_dict_reverse)
+
+
+if __name__ == '__main__':
+    test_m3_multi_devices()
+
+    print("--------------------------------")
+    print("Expected Output:")
+    print("{'colbert': [0.7798609733581543, 0.7897368669509888], 'sparse': [0.1956787109375, 0.1802978515625], 'dense': [0.6259765625, 0.67822265625], 'sparse+dense': [0.5266770720481873, 0.5633169412612915], 'colbert+sparse+dense': [0.6367570757865906, 0.6617604494094849]}")
+    print("{'colbert': [0.4524071514606476, 0.4619773030281067], 'sparse': [0.0, 0.0087890625], 'dense': [0.349853515625, 0.34765625], 'sparse+dense': [0.2691181004047394, 0.269456148147583], 'colbert+sparse+dense': [0.34880897402763367, 0.3531610071659088]}")

examples/inference/embedder/encoder_only/m3_single_device.py

@@ -0,0 +1,53 @@
+import os
+from FlagEmbedding import BGEM3FlagModel
+
+
+def test_m3_single_device():
+    model = BGEM3FlagModel(
+        'BAAI/bge-m3',
+        devices="cuda:0",   # if you don't have a GPU, you can use "cpu"
+        pooling_method='cls',
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is BGE M3?",
+        "Defination of BM25"
+    ] * 100
+    passages = [
+        "BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 
+        "BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
+    ] * 100
+    
+    queries_embeddings = model.encode_queries(
+        queries,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    passages_embeddings = model.encode_corpus(
+        passages,
+        return_dense=True,
+        return_sparse=True,
+        return_colbert_vecs=False,
+    )
+    
+    dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
+    sparse_scores = model.compute_lexical_matching_score(
+        queries_embeddings["lexical_weights"],
+        passages_embeddings["lexical_weights"],
+    )
+
+    print("Dense score:\n", dense_scores[:2, :2])
+    print("Sparse score:\n", sparse_scores[:2, :2])
+
+
+if __name__ == '__main__':
+    test_m3_single_device()
+
+    print("--------------------------------")
+    print("Expected Output:")
+    print("Dense score:")
+    print(" [[0.626  0.3477]\n [0.3496 0.678 ]]")
+    print("Sparse score:")
+    print(" [[0.19554901 0.00880432]\n [0.         0.18036556]]")

examples/inference/embedder/encoder_only/m3_single_device_compute_score.py

@@ -0,0 +1,55 @@
+import os
+from FlagEmbedding import BGEM3FlagModel
+
+
+def test_m3_single_device():
+    model = BGEM3FlagModel(
+        'BAAI/bge-m3',
+        devices="cuda:0",   # if you don't have a GPU, you can use "cpu"
+        pooling_method='cls',
+        cache_dir=os.getenv('HF_HUB_CACHE', None),
+    )
+    
+    queries = [
+        "What is BGE M3?",
+        "Defination of BM25"
+    ] * 100
+    passages = [
+        "BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 
+        "BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
+    ] * 100
+    
+    sentence_pairs = list(zip(queries, passages))
+    scores_dict = model.compute_score(
+        sentence_pairs,
+        weights_for_different_modes=[1., 0.3, 1.]
+    )
+    
+    queries.reverse()
+    sentence_pairs = list(zip(queries, passages))
+    
+    scores_dict_reverse = model.compute_score(
+        sentence_pairs,
+        weights_for_different_modes=[1., 0.3, 1.]
+    )
+    
+    scores_dict = {
+        key: value[:2]
+        for key, value in scores_dict.items()
+    }
+    scores_dict_reverse = {
+        key: value[:2]
+        for key, value in scores_dict_reverse.items()
+    }
+    
+    print(scores_dict)
+    print(scores_dict_reverse)
+
+
+if __name__ == '__main__':
+    test_m3_single_device()
+
+    print("--------------------------------")
+    print("Expected Output:")
+    print("{'colbert': [0.7798250317573547, 0.7899274826049805], 'sparse': [0.195556640625, 0.180419921875], 'dense': [0.6259765625, 0.67822265625], 'sparse+dense': [0.5266488790512085, 0.5633450746536255], 'colbert+sparse+dense': [0.6367254853248596, 0.6618592143058777]}")
+    print("{'colbert': [0.4524373412132263, 0.46213820576667786], 'sparse': [0.0, 0.0088043212890625], 'dense': [0.349609375, 0.34765625], 'sparse+dense': [0.2689302861690521, 0.26945966482162476], 'colbert+sparse+dense': [0.34871599078178406, 0.3532329499721527]}")