File size: 927 Bytes
5d50e27 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | """NVIDIA NIM embedding client."""
from openai import OpenAI
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from config import NVIDIA_API_KEY, LLM_BASE_URL, EMBED_MODEL
class Embedder:
def __init__(self):
self._client = OpenAI(base_url=LLM_BASE_URL, api_key=NVIDIA_API_KEY)
def embed_query(self, text: str) -> list[float]:
response = self._client.embeddings.create(
model=EMBED_MODEL,
input=[text],
extra_body={"input_type": "query", "truncate": "END"},
)
return response.data[0].embedding
def embed_passages(self, texts: list[str]) -> list[list[float]]:
response = self._client.embeddings.create(
model=EMBED_MODEL,
input=texts,
extra_body={"input_type": "passage", "truncate": "END"},
)
return [item.embedding for item in response.data]
|