Embed

import os
import json
import numpy as np
import cohere
from gildea import Gildea

co = cohere.ClientV2(api_key=os.environ["COHERE_API_KEY"])   # your embedder
gildea = Gildea()                                            # reads GILDEA_API_KEY

# 1. Pull the verified units for your territory; keep id + citation as provenance.
hits = gildea.search("NVIDIA data center moat durability competition", limit=30)["results"]
units = [{"id": h["unit"]["id"], "text": h["unit"]["text"], "citation": h["citation"]}
         for h in hits]

# 2. Embed the unit text with your own model -- the SAME embedder you use for your
#    own documents, so every vector is comparable in one shared space.
emb = co.embed(texts=[u["text"] for u in units], model="embed-english-v3.0",
               input_type="search_document", embedding_types=["float"])
vectors = np.array(emb.embeddings.float_)

# 3. Persist vectors next to their metadata, so you embed once and reuse across runs.
with open("embeddings.jsonl", "w") as f:
    for u, v in zip(units, vectors):
        f.write(json.dumps({"id": u["id"], "citation": u["citation"], "vector": v.tolist()}) + "\n")

# Sanity check: nearest neighbors in your own store -- the similarity a query
# across your data and Gildea's would rank on.
def near(i, k=3):
    cos = vectors @ vectors[i] / (np.linalg.norm(vectors, axis=1) * np.linalg.norm(vectors[i]))
    return [j for j in cos.argsort()[::-1] if j != i][:k]

print(f"embedded {len(units)} units -> {vectors.shape[1]}-dim, persisted to embeddings.jsonl")
print("nearest to:", units[0]["text"][:60])
for j in near(0):
    print("  ~", units[j]["text"][:60])