80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
"""Embedding generation via Ollama's qwen3-embedding:8b model."""
|
|
import hashlib
|
|
import numpy as np
|
|
|
|
|
|
def get_embedding_sync(text: str, ollama_url: str = "http://10.0.1.127:11434") -> list[float]:
|
|
"""Synchronous embedding call via Ollama /qwen3-embedding:8b."""
|
|
try:
|
|
import httpx
|
|
with httpx.Client(timeout=30) as client:
|
|
resp = client.post(
|
|
f"{ollama_url}/api/embed",
|
|
json={"model": "qwen3-embedding:8b", "input": text},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
vectors = data.get("embeddings", [])
|
|
if vectors:
|
|
# Ollama may return multiple inputs; use first
|
|
emb = vectors[0] if isinstance(vectors[0], list) else vectors
|
|
return emb
|
|
except Exception:
|
|
pass
|
|
# Fallback to deterministic feature vector if Ollama unavailable
|
|
return _hash_embedding(text)
|
|
|
|
|
|
def get_embedding(text: str, ollama_url: str = "http://10.0.1.127:11434") -> list[float]:
|
|
"""Generate an embedding — calls Ollama qwen3-embedding:8b."""
|
|
return get_embedding_sync(text, ollama_url)
|
|
|
|
|
|
async def get_embedding_async(text: str, ollama_url: str = "http://10.0.1.127:11434") -> list[float]:
|
|
"""Async embedding call via Ollama /qwen3-embedding:8b."""
|
|
try:
|
|
import httpx
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
resp = await client.post(
|
|
f"{ollama_url}/api/embed",
|
|
json={"model": "qwen3-embedding:8b", "input": text},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
vectors = data.get("embeddings", [])
|
|
if vectors:
|
|
emb = vectors[0] if isinstance(vectors[0], list) else vectors
|
|
return emb
|
|
except Exception:
|
|
pass
|
|
return _hash_embedding(text)
|
|
|
|
|
|
def _hash_embedding(text: str) -> list[float]:
|
|
"""Deterministic 4096-dim feature vector fallback (no Ollama needed)."""
|
|
feature_dim = 4096
|
|
vec = np.zeros(feature_dim, dtype=np.float32)
|
|
for n in [1, 2, 3, 4]:
|
|
tokens = [text[i:i+n] for i in range(len(text)-n+1)]
|
|
for token in tokens[:200]:
|
|
h = hashlib.md5(token.encode()).hexdigest()
|
|
for i in range(0, 12, 3):
|
|
val = (int(h[i:i+2], 16) - 128) / 128.0
|
|
feature_idx = (int(h[i+2:i+4], 16) * 37) % feature_dim
|
|
vec[feature_idx] += val
|
|
norm = np.linalg.norm(vec)
|
|
if norm > 0:
|
|
vec /= norm
|
|
return vec.tolist()
|
|
|
|
|
|
def compute_similarity(vec1: list[float], vec2: list[float]) -> float:
|
|
"""Cosine similarity between two vectors."""
|
|
v1 = np.array(vec1, dtype=np.float32)
|
|
v2 = np.array(vec2, dtype=np.float32)
|
|
if v1.shape[0] != v2.shape[0]:
|
|
return 0.0
|
|
v1 /= np.linalg.norm(v1) + 1e-8
|
|
v2 /= np.linalg.norm(v2) + 1e-8
|
|
return float(np.dot(v1, v2))
|