Acge Text Embedding

🚀 acge模型

acge模型是由合合信息技術團隊研發的一款通用文本編碼模型，屬於可變長度的向量化模型。合合信息是行業領先的人工智能及大數據科技企業，致力於通過智能文字識別及商業大數據領域的核心技術、C端和B端產品以及行業解決方案為全球企業和個人用戶提供創新的數字化、智能化服務。該模型的對外技術試用平臺為TextIn，其github開源鏈接為github。

若您有技術交流需求，請聯繫yanhui_he@intsig.net；若您有商務合作意向，請聯繫simon_liu@intsig.net。您還可以點擊圖片，掃描二維碼加入我們的微信社群。如果您想加入合合信息，從事“文檔解析”“文檔檢索”“文檔預研”相關工作，可以將簡歷投遞至min_du@intsig.net，也可直接添加HR微信詳細瞭解崗位內容。

acge模型採用了Matryoshka Representation Learning技術，如下圖所示：

建議使用的維度為1024或者1792。

✨ 主要特性

• 通用性強：作為通用的文本編碼模型，能夠處理多種文本相關任務。
• 可變長度：支持可變長度的文本向量化，具有更好的靈活性。
• 技術先進：採用了Matryoshka Representation Learning技術，提升了模型性能。

📦 安裝指南

文檔未提及安裝相關內容，故跳過此章節。

💻 使用示例

基礎用法

在sentence-transformer庫中的基本使用方法：

from sentence_transformers import SentenceTransformer

sentences = ["數據1", "數據2"]
model = SentenceTransformer('acge_text_embedding')
print(model.max_seq_length)
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

高級用法

在sentence-transformer庫中選取不同維度的使用方法：

from sklearn.preprocessing import normalize
from sentence_transformers import SentenceTransformer

sentences = ["數據1", "數據2"]
model = SentenceTransformer('acge_text_embedding')
embeddings = model.encode(sentences, normalize_embeddings=False)
matryoshka_dim = 1024
embeddings = embeddings[..., :matryoshka_dim]  # Shrink the embedding dimensions
embeddings = normalize(embeddings, norm="l2", axis=1)
print(embeddings.shape)
# => (2, 1024)

📚 詳細文檔

指標評估

C-MTEB排行榜（中文）

測試時，由於數據的隨機性、顯卡以及推理的數據類型不同，每次推理的結果會存在差異。總共進行了4次測試，使用了不同的顯卡（A10、A100）和不同的數據類型，測試結果存放在result文件夾中。最終選取了精度最低的一次測試作為最終精度測試。根據infgrad的建議，選取不同的輸入長度進行測試，當序列長度為512時測試效果最佳。

復現結果

C-MTEB：

import torch
import argparse
import functools
from C_MTEB.tasks import *
from typing import List, Dict
from sentence_transformers import SentenceTransformer
from mteb import MTEB, DRESModel


class RetrievalModel(DRESModel):
    def __init__(self, encoder, **kwargs):
        self.encoder = encoder

    def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
        input_texts = ['{}'.format(q) for q in queries]
        return self._do_encode(input_texts)

    def encode_corpus(self, corpus: List[Dict[str, str]], **kwargs) -> np.ndarray:
        input_texts = ['{} {}'.format(doc.get('title', ''), doc['text']).strip() for doc in corpus]
        input_texts = ['{}'.format(t) for t in input_texts]
        return self._do_encode(input_texts)

    @torch.no_grad()
    def _do_encode(self, input_texts: List[str]) -> np.ndarray:
        return self.encoder.encode(
            sentences=input_texts,
            batch_size=512,
            normalize_embeddings=True,
            convert_to_numpy=True
        )


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_name_or_path', default="acge_text_embedding", type=str)
    parser.add_argument('--task_type', default=None, type=str)
    parser.add_argument('--pooling_method', default='cls', type=str)
    parser.add_argument('--output_dir', default='zh_results',
                        type=str, help='output directory')
    parser.add_argument('--max_len', default=1024, type=int, help='max length')
    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
    encoder = SentenceTransformer(args.model_name_or_path).half()
    encoder.encode = functools.partial(encoder.encode, normalize_embeddings=True)
    encoder.max_seq_length = int(args.max_len)

    task_names = [t.description["name"] for t in MTEB(task_types=args.task_type,
                                                      task_langs=['zh', 'zh-CN']).tasks]
    TASKS_WITH_PROMPTS = ["T2Retrieval", "MMarcoRetrieval", "DuRetrieval", "CovidRetrieval", "CmedqaRetrieval",
                          "EcomRetrieval", "MedicalRetrieval", "VideoRetrieval"]
    for task in task_names:
        evaluation = MTEB(tasks=[task], task_langs=['zh', 'zh-CN'])
        if task in TASKS_WITH_PROMPTS:
            evaluation.run(RetrievalModel(encoder), output_folder=args.output_dir, overwrite_results=False)
        else:
            evaluation.run(encoder, output_folder=args.output_dir, overwrite_results=False)

🔧 技術細節

文檔未提及技術實現細節相關內容，故跳過此章節。

📄 許可證

文檔未提及許可證信息，故跳過此章節。