query-doc-similarity.py

# https://qwen.ai/blog?id=qwen3-vl-embedding
from pudb import set_trace
from furiosa_llm import LLM, PoolingParams

import numpy as np
import torch

queries = ["A woman playing with her dog on a beach at sunset.",
           "Pet owner training dog outdoors near water.",
           "Woman surfing on waves during a sunny day.",
           "City skyline view from a high-rise building at night.",
           ]


documents = ["A woman playing with her dog on a beach at sunrise.",
             "Dog owner training dog outdoors near ocean.",
             "Woman surfing on waves during a rainy day.",
             "City skyline view from a skyscrapers at day.",
             ]


llm = LLM("furiosa-ai/Qwen3-Embedding-8B", task='embed')

q_embeddings = llm.embed(queries)
doc_embeddings = llm.embed(documents)

q_array = torch.as_tensor([q_embeddings[i].outputs.embedding for i in range(len(queries))])
doc_array = torch.as_tensor([doc_embeddings[i].outputs.embedding for i in range(len(documents))])

# Compute similarity scores between query embeddings and document embeddings
similarity_scores = q_array @ doc_array.T
set_trace()

# Print out the similarity scores in a list format
print(similarity_scores.tolist())