75 lines
1.8 KiB
Python
75 lines
1.8 KiB
Python
"""
|
|
Embeddings module using Ollama.
|
|
|
|
Provides a simple caching layer and a function to embed text using Ollama's
|
|
embedding endpoint. No OpenAI services are used.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
from typing import List, Dict
|
|
|
|
import ollama
|
|
import numpy as np
|
|
|
|
# Cache to avoid repeated calls for the same text
|
|
_EMBED_CACHE: Dict[str, List[float]] = {}
|
|
|
|
|
|
def embed(text: str, model: str = "llama2") -> List[float]:
|
|
"""
|
|
Generate an embedding vector for the given text using Ollama.
|
|
|
|
Parameters
|
|
----------
|
|
text : str
|
|
The text to embed.
|
|
model : str, optional
|
|
The Ollama model to use for embeddings. Defaults to "llama2".
|
|
|
|
Returns
|
|
-------
|
|
List[float]
|
|
The embedding vector.
|
|
"""
|
|
if text in _EMBED_CACHE:
|
|
return _EMBED_CACHE[text]
|
|
|
|
# Ollama expects a dict with "model" and "prompt"
|
|
payload = {"model": model, "prompt": text}
|
|
try:
|
|
response = ollama.embeddings(payload)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Failed to get embeddings from Ollama: {exc}") from exc
|
|
|
|
# Ollama returns a dict with "embedding" key
|
|
embedding = response.get("embedding")
|
|
if embedding is None:
|
|
raise ValueError("Ollama response missing 'embedding' field")
|
|
|
|
_EMBED_CACHE[text] = embedding
|
|
return embedding
|
|
|
|
|
|
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
|
|
"""
|
|
Compute cosine similarity between two vectors.
|
|
|
|
Parameters
|
|
----------
|
|
vec1, vec2 : List[float]
|
|
Input vectors.
|
|
|
|
Returns
|
|
-------
|
|
float
|
|
Cosine similarity score.
|
|
"""
|
|
v1 = np.array(vec1)
|
|
v2 = np.array(vec2)
|
|
dot = np.dot(v1, v2)
|
|
norm1 = np.linalg.norm(v1)
|
|
norm2 = np.linalg.norm(v2)
|
|
if norm1 == 0 or norm2 == 0:
|
|
return 0.0
|
|
return dot / (norm1 * norm2) |