@@ -25,31 +25,33 @@ from vllm import LLM
|
||||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
|
||||
def get_detailed_instruct(task_description: str, query: str) -> str:
|
||||
return f'Instruct: {task_description}\nQuery:{query}'
|
||||
return f"Instruct: {task_description}\nQuery:{query}"
|
||||
|
||||
|
||||
def main():
|
||||
# Each query must come with a one-sentence instruction that describes the task
|
||||
task = 'Given a web search query, retrieve relevant passages that answer the query'
|
||||
task = "Given a web search query, retrieve relevant passages that answer the query"
|
||||
|
||||
queries = [
|
||||
get_detailed_instruct(task, 'What is the capital of China?'),
|
||||
get_detailed_instruct(task, 'Explain gravity')
|
||||
get_detailed_instruct(task, "What is the capital of China?"),
|
||||
get_detailed_instruct(task, "Explain gravity"),
|
||||
]
|
||||
# No need to add instruction for retrieval documents
|
||||
documents = [
|
||||
"The capital of China is Beijing.",
|
||||
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
|
||||
"Gravity is a force that attracts two bodies towards each other. "
|
||||
"It gives weight to physical objects and is responsible for the movement of planets around the sun.",
|
||||
]
|
||||
input_texts = queries + documents
|
||||
|
||||
model = LLM(model="Qwen/Qwen3-Embedding-0.6B", task="embed")
|
||||
model = LLM(model="Qwen/Qwen3-Embedding-0.6B", runner="pooling")
|
||||
|
||||
outputs = model.embed(input_texts)
|
||||
embeddings = torch.tensor([o.outputs.embedding for o in outputs])
|
||||
# Calculate the similarity scores between the first two queries and the last two documents
|
||||
scores = (embeddings[:2] @ embeddings[2:].T)
|
||||
scores = embeddings[:2] @ embeddings[2:].T
|
||||
print(scores.tolist())
|
||||
# [[0.7620252966880798, 0.14078938961029053], [0.1358368694782257, 0.6013815999031067]]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user