# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "requests"]
# ///

# dense-retrieve.py
import numpy as np, requests

QUERY_PREFIX, DOC_PREFIX = "query: ", "document: "

def embed(text: str) -> np.ndarray:
    r = requests.post(
        "http://localhost:8080/v1/embeddings",
        json={"input": text},
    )
    v = np.array(r.json()["data"][0]["embedding"])
    return v / np.linalg.norm(v)

docs = [
    "hi",
    "it is a bear",
    "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.",
]
query = "What is panda?"

q = embed(QUERY_PREFIX + query)
for doc in docs:
    d = embed(DOC_PREFIX + doc)
    print(f"Score: {float(q @ d):.4f} | Q: {query} | D: {doc}")