embeddinggemma-300m-onnx

模型概述

embeddinggemma-300m-onnx 是 Google Gemma 嵌入模型的高效 ONNX 导出版本，包含约 3 亿参数。
它能为文本生成高质量的语义嵌入，适用于各种 NLP 任务，包括句子相似度、文本聚类、分类和检索。

将原始 Gemma 嵌入模型转换为 ONNX 格式，可以通过 ONNX Runtime 在 CPU 和 GPU 上实现硬件无关的优化推理。

原始模型参考

该模型基于以高效和多语言能力著称的 Google Gemma 嵌入架构。
请参考原始 Hugging Face 仓库和文档：
google/embeddinggemma-300m
在研究或生产中使用此模型时，请引用原始论文。

预期用途

• 语义相似度评分
• 搜索和推荐系统的嵌入生成
• 文本分类和聚类
• 需要 ONNX 兼容模型的可扩展、低延迟推理场景

仓库文件

安装

pip install onnxruntime transformers huggingface_hub

使用方法

1. 克隆或下载文件后在本地使用模型

import onnxruntime as ort
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(".")
session = ort.InferenceSession("./model.onnx")

text = "Example input text"
inputs = tokenizer(text, return_tensors="np")
outputs = session.run(None, dict(inputs))
embeddings = outputs
print(embeddings)

2. 直接从 Hugging Face Hub 使用模型（无需手动下载）

from huggingface_hub import hf_hub_download
import onnxruntime as ort
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("be1newinner/embeddinggemma-300m-onnx")
onnx_model_path = hf_hub_download(
    repo_id="be1newinner/embeddinggemma-300m-onnx",
    filename="model.onnx"
)
session = ort.InferenceSession(onnx_model_path)

text = "Example input text"
inputs = tokenizer(text, return_tensors="np")
outputs = session.run(None, dict(inputs))
embeddings = outputs
print(embeddings)

3. 生成固定 768 维嵌入向量作为输出

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download

class GemmaEmbedder:
    """
    A class to generate embeddings using a Gemma model in ONNX format.
    """

    def __init__(self, model_repo="be1newinner/embeddinggemma-300m-onnx"):
        """
        Initializes the GemmaEmbedder by loading the tokenizer and ONNX model.

        Args:
            model_repo (str): The repository ID of the ONNX model on Hugging Face Hub.
        """
        # Load the tokenizer from the Hugging Face Hub
        self.tokenizer = AutoTokenizer.from_pretrained(model_repo)

        # Download and load the ONNX model
        onnx_model_path = hf_hub_download(repo_id=model_repo, filename="model.onnx")
        self.session = ort.InferenceSession(onnx_model_path)

    def generate(self, text: str):
        """
        Generates a fixed-size embedding for the input text.

        Args:
            text (str): The input text to embed.

        Returns:
            np.ndarray: The generated embedding as a NumPy array.
        """
        # Tokenize the input text, padding and truncating to a consistent length
        inputs = self.tokenizer(
            text, return_tensors="np", padding=True, truncation=True
        )

        # Run the ONNX model to get the last hidden states
        outputs = self.session.run(None, dict(inputs))

        # Perform mean pooling to get a fixed-size embedding
        last_hidden_states = outputs[0]
        input_mask_expanded = np.expand_dims(inputs["attention_mask"], -1).astype(float)
        sum_embeddings = np.sum(last_hidden_states * input_mask_expanded, 1)
        sum_mask = np.clip(input_mask_expanded.sum(1), a_min=1e-9, a_max=None)

        pooled_embeddings = sum_embeddings / sum_mask
        return pooled_embeddings


# Create a global instance of the embedder to avoid reloading the model
embedder = GemmaEmbedder()


def generate(text: str):
    """
    A convenience function to generate embeddings using the global embedder instance.

    Args:
        text (str): The input text to embed.

    Returns:
        np.ndarray: The generated embedding.
    """
    return embedder.generate(text)


if __name__ == "__main__":
    # Example usage of the generate function
    embeddings = generate("Example input text")
    print(embeddings)
    print(f"Embedding shape: {embeddings.shape}")

引用

如果您在工作中使用此模型，请引用：

• 原始 Google Gemma 嵌入模型论文（链接见原始仓库）。
• 本仓库及 Hugging Face 托管（如适用）。

欢迎通过本仓库的 GitHub 页面贡献改进或报告问题。

最后更新：2025 年 11 月
维护者：be1newinner