diff --git a/ragengine/__init__.py b/ragengine/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/ragengine/embedding/__init__.py b/ragengine/embedding/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/ragengine/embedding/base.py b/ragengine/embedding/base.py new file mode 100644 index 000000000..a1c371937 --- /dev/null +++ b/ragengine/embedding/base.py @@ -0,0 +1,13 @@ +from abc import ABC, abstractmethod + + +class BaseEmbeddingModel(ABC): + @abstractmethod + def get_text_embedding(self, text: str): + """Returns the text embedding for a given input string.""" + pass + + @abstractmethod + def get_embedding_dimension(self) -> int: + """Returns the embedding dimension for the model.""" + pass \ No newline at end of file diff --git a/ragengine/embedding/huggingface_local.py b/ragengine/embedding/huggingface_local.py new file mode 100644 index 000000000..3dab51e9a --- /dev/null +++ b/ragengine/embedding/huggingface_local.py @@ -0,0 +1,19 @@ +from llama_index.embeddings.huggingface import HuggingFaceEmbedding + +from .base import BaseEmbeddingModel + + +class LocalHuggingFaceEmbedding(BaseEmbeddingModel): + def __init__(self, model_name: str): + self.model = HuggingFaceEmbedding(model_name=model_name) # TODO: Ensure/test loads on GPU (when available) + + def get_text_embedding(self, text: str): + """Returns the text embedding for a given input string.""" + return self.model.get_text_embedding(text) + + def get_embedding_dimension(self) -> int: + """Infers the embedding dimension by making a local call to get the embedding of a dummy text.""" + dummy_input = "This is a dummy sentence." + embedding = self.get_text_embedding(dummy_input) + + return len(embedding) \ No newline at end of file diff --git a/ragengine/embedding/huggingface_remote.py b/ragengine/embedding/huggingface_remote.py new file mode 100644 index 000000000..0f8e79181 --- /dev/null +++ b/ragengine/embedding/huggingface_remote.py @@ -0,0 +1,20 @@ +from llama_index.embeddings.huggingface_api import \ + HuggingFaceInferenceAPIEmbedding + +from .base import BaseEmbeddingModel + + +class RemoteHuggingFaceEmbedding(BaseEmbeddingModel): + def __init__(self, model_name: str, api_key: str): + self.model = HuggingFaceInferenceAPIEmbedding(model_name=model_name, token=api_key) + + def get_text_embedding(self, text: str): + """Returns the text embedding for a given input string.""" + return self.model.get_text_embedding(text) + + def get_embedding_dimension(self) -> int: + """Infers the embedding dimension by making a remote call to get the embedding of a dummy text.""" + dummy_input = "This is a dummy sentence." + embedding = self.get_text_embedding(dummy_input) + + return len(embedding) diff --git a/ragengine/requirements.txt b/ragengine/requirements.txt new file mode 100644 index 000000000..18ca062fc --- /dev/null +++ b/ragengine/requirements.txt @@ -0,0 +1,7 @@ +# RAG Library Requirements +llama-index +llama-index-embeddings-huggingface +fastapi +faiss-cpu +llama-index-vector-stores-faiss +uvicorn