Spaces:
Runtime error
Runtime error
File size: 1,946 Bytes
105b369 |
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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from typing import Optional, Dict, List, Tuple, Any
from phi.embedder.base import Embedder
from phi.utils.log import logger
try:
from ollama import Client as OllamaClient
except ImportError:
logger.error("`ollama` not installed")
raise
class OllamaEmbedder(Embedder):
model: str = "openhermes"
dimensions: int = 4096
host: Optional[str] = None
timeout: Optional[Any] = None
options: Optional[Any] = None
client_kwargs: Optional[Dict[str, Any]] = None
ollama_client: Optional[OllamaClient] = None
@property
def client(self) -> OllamaClient:
if self.ollama_client:
return self.ollama_client
_ollama_params: Dict[str, Any] = {}
if self.host:
_ollama_params["host"] = self.host
if self.timeout:
_ollama_params["timeout"] = self.timeout
if self.client_kwargs:
_ollama_params.update(self.client_kwargs)
return OllamaClient(**_ollama_params)
def _response(self, text: str) -> Dict[str, Any]:
kwargs: Dict[str, Any] = {}
if self.options is not None:
kwargs["options"] = self.options
return self.client.embeddings(prompt=text, model=self.model, **kwargs) # type: ignore
def get_embedding(self, text: str) -> List[float]:
try:
response = self._response(text=text)
if response is None:
return []
return response.get("embedding", [])
except Exception as e:
logger.warning(e)
return []
def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
embedding = []
usage = None
try:
response = self._response(text=text)
if response is not None:
embedding = response.get("embedding", [])
except Exception as e:
logger.warning(e)
return embedding, usage
|