|
|
|
|
|
|
|
|
def __init__(self, key, model_name, **kwargs): |
|
|
def __init__(self, key, model_name, **kwargs): |
|
|
self.client = Client(host=kwargs["base_url"]) if not key or key == "x" else Client(host=kwargs["base_url"], headers={"Authorization": f"Bearer {key}"}) |
|
|
self.client = Client(host=kwargs["base_url"]) if not key or key == "x" else Client(host=kwargs["base_url"], headers={"Authorization": f"Bearer {key}"}) |
|
|
self.model_name = model_name |
|
|
self.model_name = model_name |
|
|
|
|
|
self.keep_alive = kwargs.get("ollama_keep_alive", int(os.environ.get("OLLAMA_KEEP_ALIVE", -1))) |
|
|
|
|
|
|
|
|
def encode(self, texts: list): |
|
|
def encode(self, texts: list): |
|
|
arr = [] |
|
|
arr = [] |
|
|
|
|
|
|
|
|
# remove special tokens if they exist |
|
|
# remove special tokens if they exist |
|
|
for token in OllamaEmbed._special_tokens: |
|
|
for token in OllamaEmbed._special_tokens: |
|
|
txt = txt.replace(token, "") |
|
|
txt = txt.replace(token, "") |
|
|
res = self.client.embeddings(prompt=txt, model=self.model_name, options={"use_mmap": True}, keep_alive=-1) |
|
|
|
|
|
|
|
|
res = self.client.embeddings(prompt=txt, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive) |
|
|
try: |
|
|
try: |
|
|
arr.append(res["embedding"]) |
|
|
arr.append(res["embedding"]) |
|
|
except Exception as _e: |
|
|
except Exception as _e: |
|
|
|
|
|
|
|
|
# remove special tokens if they exist |
|
|
# remove special tokens if they exist |
|
|
for token in OllamaEmbed._special_tokens: |
|
|
for token in OllamaEmbed._special_tokens: |
|
|
text = text.replace(token, "") |
|
|
text = text.replace(token, "") |
|
|
res = self.client.embeddings(prompt=text, model=self.model_name, options={"use_mmap": True}, keep_alive=-1) |
|
|
|
|
|
|
|
|
res = self.client.embeddings(prompt=text, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive) |
|
|
try: |
|
|
try: |
|
|
return np.array(res["embedding"]), 128 |
|
|
return np.array(res["embedding"]), 128 |
|
|
except Exception as _e: |
|
|
except Exception as _e: |