diff --git a/core/cat/factory/dumb_embedder.py b/core/cat/factory/custom_embedder.py similarity index 72% rename from core/cat/factory/dumb_embedder.py rename to core/cat/factory/custom_embedder.py index 825c6474..e54f7452 100644 --- a/core/cat/factory/dumb_embedder.py +++ b/core/cat/factory/custom_embedder.py @@ -1,8 +1,11 @@ +import os import string +import json from typing import List from itertools import combinations from sklearn.feature_extraction.text import CountVectorizer from langchain.embeddings.base import Embeddings +import httpx class DumbEmbedder(Embeddings): @@ -41,3 +44,21 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]: def embed_query(self, text: str) -> List[float]: """Embed a string of text and returns the embedding vector as a list of floats.""" return self.embedder.transform([text]).astype(float).todense().tolist()[0] + + +class CustomOpenAIEmbeddings(Embeddings): + def __init__(self, url): + self.url = os.path.join(url, "v1/embeddings") + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + payload = json.dumps({"input": texts}) + ret = httpx.post(self.url, data=payload, timeout=None) + ret.raise_for_status() + return [e['embedding'] for e in ret.json()['data']] + + def embed_query(self, text: str) -> List[float]: + payload = json.dumps({"input": text}) + ret = httpx.post(self.url, data=payload, timeout=None) + ret.raise_for_status() + return ret.json()['data'][0]['embedding'] + \ No newline at end of file diff --git a/core/cat/factory/custom_llm.py b/core/cat/factory/custom_llm.py index 58a8b607..2a2d37b5 100644 --- a/core/cat/factory/custom_llm.py +++ b/core/cat/factory/custom_llm.py @@ -66,10 +66,12 @@ def _identifying_params(self) -> Mapping[str, Any]: class CustomOpenAI(OpenAI): + url: str + def __init__(self, **kwargs): - model_kwargs = { - 'repeat_penalty': kwargs.pop('repeat_penalty') + 'repeat_penalty': kwargs.pop('repeat_penalty'), + 'top_k': kwargs.pop('top_k') } stop = kwargs.pop('stop', None) @@ -81,4 +83,7 @@ def __init__(self, **kwargs): model_kwargs=model_kwargs, **kwargs ) - self.openai_api_base = os.path.join(kwargs['url'], "v1") \ No newline at end of file + + self.url = kwargs['url'] + self.openai_api_base = os.path.join(self.url, "v1") + \ No newline at end of file diff --git a/core/cat/factory/embedder.py b/core/cat/factory/embedder.py index 5df0030c..464a1f68 100644 --- a/core/cat/factory/embedder.py +++ b/core/cat/factory/embedder.py @@ -1,7 +1,7 @@ import langchain from pydantic import PyObject, BaseSettings -from cat.factory.dumb_embedder import DumbEmbedder +from cat.factory.custom_embedder import DumbEmbedder, CustomOpenAIEmbeddings # Base class to manage LLM configuration. @@ -36,11 +36,22 @@ class EmbedderDumbConfig(EmbedderSettings): class Config: schema_extra = { - "name_human_readable": "Dumb Embedder", + "humanReadableName": "Dumb Embedder", "description": "Configuration for default embedder. It encodes the pairs of characters", } +class EmbedderLlamaCppConfig(EmbedderSettings): + url: str + _pyclass = PyObject = CustomOpenAIEmbeddings + + class Config: + schema_extra = { + "humanReadableName": "Self-hosted llama-cpp-python embedder", + "description": "Self-hosted llama-cpp-python embedder", + } + + class EmbedderOpenAIConfig(EmbedderSettings): openai_api_key: str model: str = "text-embedding-ada-002" @@ -98,6 +109,7 @@ class Config: SUPPORTED_EMDEDDING_MODELS = [ EmbedderDumbConfig, EmbedderFakeConfig, + EmbedderLlamaCppConfig, EmbedderOpenAIConfig, EmbedderAzureOpenAIConfig, EmbedderCohereConfig, diff --git a/core/cat/mad_hatter/core_plugin/hooks/models.py b/core/cat/mad_hatter/core_plugin/hooks/models.py index a2779124..a41a4e4c 100644 --- a/core/cat/mad_hatter/core_plugin/hooks/models.py +++ b/core/cat/mad_hatter/core_plugin/hooks/models.py @@ -15,6 +15,7 @@ from langchain import HuggingFaceHub from langchain.chat_models import AzureChatOpenAI from cat.mad_hatter.decorators import hook +from cat.factory.custom_llm import CustomOpenAI @hook(priority=0) @@ -142,6 +143,14 @@ def get_language_embedder(cat) -> embedders.EmbedderSettings: } ) + # Llama-cpp-python + elif type(cat._llm) in [CustomOpenAI]: + embedder = embedders.EmbedderLlamaCppConfig.get_embedder_from_config( + { + "url": cat._llm.url + } + ) + else: # If no embedder matches vendor, and no external embedder is configured, we use the DumbEmbedder. # `This embedder is not a model properly trained