| """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] |
|
|
