File size: 1,130 Bytes
ad82189 | 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 30 | """Supply zembed's saved prompt and terminal suffix to Infinity's /embed API."""
import json
import urllib.error
import urllib.request
def embed(texts, kind="document", base_url="http://localhost:7997", model="michaelfeil/zembed-1-infinity"):
if isinstance(texts, str):
texts = [texts]
texts = list(texts)
if kind not in ("query", "document") or not texts or any(not isinstance(text, str) for text in texts):
raise ValueError("Provide text strings and kind='query' or kind='document'")
body = {
"model": model,
"inputs": [text + "<|im_end|>\n" for text in texts],
"prompt_name": kind,
"normalize": True,
"truncate": False,
}
request = urllib.request.Request(
base_url.rstrip("/") + "/embed",
data=json.dumps(body).encode(),
headers={"Content-Type": "application/json"},
)
try:
with urllib.request.urlopen(request, timeout=300) as response:
return json.load(response)
except urllib.error.HTTPError as error:
raise RuntimeError(f"HTTP {error.code}: {error.read().decode()}") from error
|