mlx embeddings
1.0.0
MLX-Embeddings は、MLX を使用して Mac 上でビジョンおよび言語埋め込みモデルをローカルに実行するためのパッケージです。
pip を使用して mlx-embeddings をインストールできます。
pip install mlx-embeddings
単一のテキストの埋め込みを生成するには:
import mlx . core as mx
from mlx_embeddings . utils import load
# Load the model and tokenizer
model , tokenizer = load ( "sentence-transformers/all-MiniLM-L6-v2" )
# Prepare the text
text = "I like reading"
# Tokenize and generate embedding
input_ids = tokenizer . encode ( text , return_tensors = "mlx" )
outputs = model ( input_ids )
embeddings = outputs [ 0 ][:, 0 , :]
埋め込みを使用して複数のテキストを比較するには:
from sklearn . metrics . pairwise import cosine_similarity
import matplotlib . pyplot as plt
import seaborn as sns
import mlx . core as mx
from mlx_embeddings . utils import load
# Load the model and tokenizer
model , tokenizer = load ( "sentence-transformers/all-MiniLM-L6-v2" )
def get_embedding ( text , model , tokenizer ):
input_ids = tokenizer . encode ( text , return_tensors = "mlx" , padding = True , truncation = True , max_length = 512 )
outputs = model ( input_ids )
embeddings = outputs [ 0 ][:, 0 , :][ 0 ]
return embeddings
# Sample texts
texts = [
"I like grapes" ,
"I like fruits" ,
"The slow green turtle crawls under the busy ant."
]
# Generate embeddings
embeddings = [ get_embedding ( text , model , tokenizer ) for text in texts ]
# Compute similarity
similarity_matrix = cosine_similarity ( embeddings )
# Visualize results
def plot_similarity_matrix ( similarity_matrix , labels ):
plt . figure ( figsize = ( 5 , 4 ))
sns . heatmap ( similarity_matrix , annot = True , cmap = 'coolwarm' , xticklabels = labels , yticklabels = labels )
plt . title ( 'Similarity Matrix Heatmap' )
plt . tight_layout ()
plt . show ()
labels = [ f"Text { i + 1 } " for i in range ( len ( texts ))]
plot_similarity_matrix ( similarity_matrix , labels )
複数のテキストを一度に処理するには:
from sklearn . metrics . pairwise import cosine_similarity
import matplotlib . pyplot as plt
import seaborn as sns
import mlx . core as mx
from mlx_embeddings . utils import load
# Load the model and tokenizer
model , tokenizer = load ( "sentence-transformers/all-MiniLM-L6-v2" )
def get_embedding ( texts , model , tokenizer ):
inputs = tokenizer . batch_encode_plus ( texts , return_tensors = "mlx" , padding = True , truncation = True , max_length = 512 )
outputs = model (
inputs [ "input_ids" ],
attention_mask = inputs [ "attention_mask" ]
)
return outputs [ 0 ]
def compute_and_print_similarity ( embeddings ):
B , Seq_len , dim = embeddings . shape
embeddings_2d = embeddings . reshape ( B , - 1 )
similarity_matrix = cosine_similarity ( embeddings_2d )
print ( "Similarity matrix between sequences:" )
print ( similarity_matrix )
print ( " n " )
for i in range ( B ):
for j in range ( i + 1 , B ):
print ( f"Similarity between sequence { i + 1 } and sequence { j + 1 } : { similarity_matrix [ i ][ j ]:.4f } " )
return similarity_matrix
# Sample texts
texts = [
"I like grapes" ,
"I like fruits" ,
"The slow green turtle crawls under the busy ant."
]
embeddings = get_embedding ( texts , model , tokenizer )
similarity_matrix = compute_and_print_similarity ( embeddings )
# Visualize results
labels = [ f"Text { i + 1 } " for i in range ( len ( texts ))]
plot_similarity_matrix ( similarity_matrix , labels )
MLX-Embeddings は、テキスト埋め込みタスク用のさまざまなモデル アーキテクチャをサポートしています。現在サポートされているアーキテクチャの内訳は次のとおりです。
私たちは追加のモデル アーキテクチャのサポートを拡大するために継続的に取り組んでいます。サポートされているモデルとその特定のバージョンの最新リストについては、GitHub リポジトリまたはドキュメントを確認してください。
MLX-Embeddings への貢献は大歓迎です。詳細については、投稿ガイドラインを参照してください。
このプロジェクトは、GNU General Public License v3 に基づいてライセンスされています。
ご質問や問題がある場合は、GitHub リポジトリで問題を開いてください。