Gte Small

🚀 gte-small

General Text Embeddings (GTE) モデルは、Alibaba DAMO Academyによって開発されたテキスト埋め込みモデルです。BERTフレームワークをベースに構築され、大規模な関連性のあるテキストペアのコーパスで訓練されています。これにより、様々なドメインやシナリオに対応し、情報検索、意味的なテキストの類似性評価、テキストの再ランキングなど、様々な下流タスクに適用できます。

✨ 主な機能

• 大規模コーパスで訓練された、汎用的なテキスト埋め込みモデル。
• 様々な下流タスクに適用可能。
• PythonとJavaScriptの両方で使用可能。

📦 インストール

このモデルは、transformersライブラリを介して使用できます。Python環境では、以下のコマンドでインストールできます。

pip install transformers torch

JavaScript環境では、@xenova/transformersを使用します。

npm install @xenova/transformers

💻 使用例

基本的な使用法

Pythonでの使用

import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]

tokenizer = AutoTokenizer.from_pretrained("Supabase/gte-small")
model = AutoModel.from_pretrained("Supabase/gte-small")

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# (Optionally) normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())

JavaScriptでの使用

DenoまたはSupabase Edge Functionsでの使用例：

import { serve } from 'https://deno.land/std@0.168.0/http/server.ts'
import { env, pipeline } from 'https://cdn.jsdelivr.net/npm/@xenova/transformers@2.5.0'

// Configuration for Deno runtime
env.useBrowserCache = false;
env.allowLocalModels = false;

const pipe = await pipeline(
  'feature-extraction',
  'Supabase/gte-small',
);

serve(async (req) => {
  // Extract input string from JSON body
  const { input } = await req.json();

  // Generate the embedding from the user input
  const output = await pipe(input, {
    pooling: 'mean',
    normalize: true,
  });

  // Extract the embedding output
  const embedding = Array.from(output.data);

  // Return the embedding
  return new Response(
    JSON.stringify({ embedding }),
    { headers: { 'Content-Type': 'application/json' } }
  );
});

高度な使用法

Pythonでのsentence-transformersを使用した例

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

sentences = ['That is a happy person', 'That is a very happy person']

model = SentenceTransformer('Supabase/gte-small')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

ブラウザ内でのJavaScriptの使用例

<script type="module">

import { pipeline } from 'https://cdn.jsdelivr.net/npm/@xenova/transformers@2.5.0';

const pipe = await pipeline(
  'feature-extraction',
  'Supabase/gte-small',
);

// Generate the embedding from text
const output = await pipe('Hello world', {
  pooling: 'mean',
  normalize: true,
});

// Extract the embedding output
const embedding = Array.from(output.data);

console.log(embedding);

</script>

Node.jsまたはWebpackでの使用例

import { pipeline } from '@xenova/transformers';

const pipe = await pipeline(
  'feature-extraction',
  'Supabase/gte-small',
);

// Generate the embedding from text
const output = await pipe('Hello world', {
  pooling: 'mean',
  normalize: true,
});

// Extract the embedding output
const embedding = Array.from(output.data);

console.log(embedding);

📚 ドキュメント

メトリクス

GTEモデルの性能は、MTEBベンチマークで他の人気のテキスト埋め込みモデルと比較されています。詳細な比較結果については、MTEBリーダーボードを参照してください。

制限事項

このモデルは英語のテキストのみを対象としており、長いテキストは最大512トークンに切り捨てられます。

📄 ライセンス

このモデルはMITライセンスの下で提供されています。