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