使用 Sentence Transformers 產生嵌入¶

from sentence_transformers import SentenceTransformer

model_id = "google/embeddinggemma-2"
model = SentenceTransformer(model_id)

print(f"device: {model.device}")
print(f"dtype: {model.dtype}")
print(model)
print("Total number of parameters in the model:", sum(p.numel() for p in model.parameters()))

words = ["king", "queen", "car"]

# Generate embeddings using model.encode()
embeddings = model.encode(words)

print(embeddings)
for idx, embedding in enumerate(embeddings):
    print(f"Embedding {idx+1} (shape): {embedding.shape}")

sentence_high = [
    "The chef prepared a delicious meal for the guests.",
    "A tasty dinner was cooked by the chef for the visitors.",
]
sentence_medium = [
    "She is an expert in machine learning.",
    "He has a deep interest in artificial intelligence.",
]
sentence_low = [
    "The weather in Tokyo is sunny today.",
    "I need to buy groceries for the week.",
]

for pair in [sentence_high, sentence_medium, sentence_low]:
    print("🙋‍♂️", pair)
    pair_embeddings = model.encode(pair)
    score = model.similarity(pair_embeddings[0], pair_embeddings[1]).item()
    print(f" `-> 🤖 similarity score: {score:.4f}")

print("Available task prompts:")
for name, prefix in model.prompts.items():
    print(f"  {name}: {prefix!r}")
print("-" * 80)

for pair in [sentence_high, sentence_medium, sentence_low]:
    print("🙋‍♂️", pair)
    pair_embeddings = model.encode(pair, prompt_name="STS")
    score = model.similarity(pair_embeddings[0], pair_embeddings[1]).item()
    print(f" `-> 🤖 STS score: {score:.4f}")

query = "How do I reduce memory usage when running EmbeddingGemma 2?"

documents = [
    {
        "title": "Model Initialization and Memory Optimization",
        "text": "You can disable the vision or audio encoders by passing vision_config=None or audio_config=None inside config_kwargs.",
    },
    {
        "title": "Matryoshka Representation Learning",
        "text": "MRL allows truncating 768-dimensional embeddings to 512 or 256 dimensions for faster vector search and smaller indexes.",
    },
    {
        "title": "Audio Processing Specifications",
        "text": "Audio waveforms are automatically resampled to 16 kHz before feature extraction.",
    },
]

# Encode query with Retrieval-query prompt
query_emb = model.encode(query, prompt_name="Retrieval-query")

# Encode documents with their titles embedded in the prompt
doc_embeddings = [
    model.encode(doc["text"], prompt=f"title: {doc['title']} | text: ")
    for doc in documents
]

scores = model.similarity(query_emb, doc_embeddings)[0]
for doc, score in zip(documents, scores):
    print(f"[Score: {score:.4f}] {doc['title']}")

labels = ["Billing Issue", "Technical Support", "Sales Inquiry"]

tickets = [
    "Excuse me, the app freezes on the login screen. It won't work even when I try to reset my password.",
    "I would like to inquire about your enterprise plan pricing and features for a team of 50 people.",
]

# Encode labels and incoming tickets using the Clustering prompt
label_embeddings = model.encode(labels, prompt_name="Clustering")
ticket_embeddings = model.encode(tickets, prompt_name="Clustering")

# Compute similarity matrix and find the best-matching label for each ticket
similarities = model.similarity(ticket_embeddings, label_embeddings)
print("Similarity matrix:\n", similarities)

predicted_indices = similarities.argmax(dim=1)

for ticket, pred_idx in zip(tickets, predicted_indices):
    print(f"\n🙋‍♂️ Ticket: {ticket}")
    print(f" `-> 🤖 Predicted Category: {labels[pred_idx]}")

labels = ["Positive", "Neutral", "Negative"]

sentences = [
    # Positive
    "The sunset over the ridge was absolutely breathtaking.",
    "This is hands down the best coffee I have ever tasted.",
    # Neutral
    "The package arrived on Tuesday afternoon.",
    "The meeting is scheduled for 10:00 AM in Conference Room B.",
    # Negative
    "The battery dies after barely two hours of use.",
    "It felt like an utter waste of time and money.",
]

# Encode labels and incoming sentences using the Classification prompt
label_embeddings = model.encode(labels, prompt_name="Classification")
sentence_embeddings = model.encode(sentences, prompt_name="Classification")

# Compute similarity matrix and find the best-matching label for each sentence
similarities = model.similarity(sentence_embeddings, label_embeddings)
print("Similarity matrix:\n", similarities)

predicted_indices = similarities.argmax(dim=1)

for sentence, pred_idx in zip(sentences, predicted_indices):
    print(f"\n🙋‍♂️ Sentence: {sentence}")
    print(f" `-> 🤖 Predicted Sentiment: {labels[pred_idx]}")

words = ["king", "queen", "car"]

def evaluate_mrl(embeddings, sim_fn_name):
    print(f"Embedding shape: {embeddings.shape} | Similarity function: {sim_fn_name}")
    sims = model.similarity(embeddings[0], embeddings[1:])
    for idx, word in enumerate(words[1:]):
        print(f"  🙋‍♂️ 'king' vs. '{word}' -> 🤖 score: {sims[0][idx].item():.4f}")

# 1. Full 768-dimensional embeddings
emb_768 = model.encode(words, prompt_name="STS", normalize_embeddings=True)
evaluate_mrl(emb_768, model.similarity_fn_name)

# 2. Truncate to 512 dimensions at inference time
emb_512 = model.encode(
    words,
    prompt_name="STS",
    truncate_dim=512,
    normalize_embeddings=True,
)
evaluate_mrl(emb_512, model.similarity_fn_name)

# 3. Initialize model with default 128-dim truncation and dot-product similarity
fast_model = SentenceTransformer(
    model_id,
    truncate_dim=128,
    similarity_fn_name="dot",
    device=model.device,
)

emb_128 = fast_model.encode(words, prompt_name="STS", normalize_embeddings=True)
evaluate_mrl(emb_128, fast_model.similarity_fn_name)