vLLM Example

from vllm import LLM
import torch
import torch.nn.functional as F
import numpy as np

def cosine_similarity(x, y):
    """Compute cosine similarity between two tensors."""
    x = F.normalize(x, p=2, dim=1)
    y = F.normalize(y, p=2, dim=1)
    return x @ y.T

llm = LLM(
    model='jinaai/jina-embeddings-v5-text-0.6B-pair-ckpt',
    task="embed",
)
print("Model loaded successfully!")

# Example texts to encode
example_texts = [
    "The quantum entanglement of particles across vast cosmic distances reveals the profound interconnectedness of the universe.",
    "In the depths of the ocean, bioluminescent creatures create their own constellations in an eternal midnight sky.",
    "A neural network dreams in patterns, finding beauty in the chaos of data streams and learning to see the world through numbers.",
    "The ancient library contained scrolls written in languages forgotten by time, each page holding secrets of civilizations long vanished.",
    "Through the telescope, galaxies spiral like cosmic whirlpools, each one a universe of possibilities waiting to be explored.",
]

print(f"\nEncoding {len(example_texts)} texts...")

# Encode with vLLM's embed method (uses last token pooling)
outputs = llm.embed(example_texts)

# Collect embeddings into a single tensor
emb_list = []
for i, out in enumerate(outputs):
    vec = out.outputs.embedding
    vec = torch.tensor(vec)        
    emb_list.append(vec)
embeddings = torch.stack(emb_list, dim=0)

# Convert to numpy for analysis
embeddings_np = embeddings.cpu().numpy()

# Print results
print(f"\nEmbeddings shape: {embeddings_np.shape}")
print(f"Embedding dimension: {embeddings_np.shape[1]}")

# Compute cosine similarity matrix
similarity_matrix = cosine_similarity(embeddings, embeddings)
print(f"\nCosine similarity matrix:")
print(similarity_matrix.cpu().numpy())