Tarka-Embedding-150M-V1 (ONNX)

Tarka-AIR/Tarka-Embedding-150M-V1 的 ONNX 版本。

• 嵌入维度: 768
• 上下文长度: 2048
• 模型大小: ~600MB

使用方法

使用 ONNX Runtime (Python)

import onnxruntime as ort
from transformers import AutoTokenizer

session = ort.InferenceSession("tarka-150m-v1-onnx/model.onnx")
tokenizer = AutoTokenizer.from_pretrained("permutans/Tarka-Embedding-150M-V1-ONNX")

texts = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium.",
]

embeddings = []
for text in texts:
    inputs = tokenizer(text, return_tensors="np")
    _, sentence_embedding = session.run(
        None,
        {"input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"]},
    )
    embeddings.append(sentence_embedding[0])

import numpy as np

embeddings = np.array(embeddings)
print(embeddings.shape)  # (3, 768)

# Compute cosine similarities
from sklearn.metrics.pairwise import cosine_similarity

similarities = cosine_similarity(embeddings)
print(similarities)

使用 FastEmbed (Rust)

兼容 fastembed-rs，用于高性能嵌入生成。

模型输出

• token_embeddings: 词元级嵌入 (batch_size, sequence_length, 768)
• sentence_embedding: 池化后的句子嵌入 (batch_size, 768) - 大多数任务请使用此项

性能

此 ONNX 导出版本同时支持 CPU 和 CUDA 执行提供程序，便于灵活部署。