chroma-core / chroma-core/chroma
[Feature Request]: Llama_Cpp_Python Support for embedding function
- Dominant language
- Rust
- Stars
- 29.3k
- Forks
- 2.5k
- Avg merge
- 1d 4h
- Merged PRs (30d)
- 38
Description
### Describe the problem
I would like if the collections class had a built-in function for using a local model and llamacpp to embed the documents
### Describe the proposed solution
Since I have no idea how to contribute to a open repo, I wish this was a function in the collection embedded functions list:
```python
from chromadb import Documents, EmbeddingFunction, Embeddings
from llama_cpp import Llama
from torch import cuda
class LlamaCppEmbeddingFunction(EmbeddingFunction):
def __init__(self, model_path: str, **kwargs: Any):
"""
Initialize the LlamaCppEmbeddingFunction. This function will embed documents using the Llama-CPP-Python library.
Args:
model_path (str): Path to the model file.
kwargs: Additional arguments to pass to the Llama constructor.
* n_ctx (int): The context size.
* n_threads (int): The number of cpu threads to use.
* n_gpu_layers (int): The number of layers to run on the GPU.
"""
self.model_path = model_path
# Check if verbose is in kwargs, if not set to False
if 'verbose' not in kwargs:
kwargs['verbose'] = False
# Force embedding to be True
kwargs['embedding'] = True
# Check if the computer has a GPU, if not set n_gpu_layers to 0
if cuda.is_available():
if 'n_gpu_layers' not in kwargs:
kwargs['n_gpu_layers'] = 1
else:
kwargs['n_gpu_layers'] = 0
try:
self.llm_embedding = Llama(model_path, **kwargs)
except Exception as e:
raise Exception(f"Error initializing LlamaCppEmbeddingFunction: {e}")
def __call__(self, input: Documents) -> Embeddings:
# Create embeddings
llama_embeddings = [embedding['embedding'] for embedding in self.llm_embedding.create_embedding(list(input))['data']]
# Convert to numpy array
llama_embeddings = np.array(llama_embeddings)
# embed the documents somehow
return cast(
Embeddings,
llama_embeddings.tolist()
)
```
### Alternatives considered
_No response_
### Importance
i cannot use Chroma without it
### Additional Information
The hope is to use it for my vector database like so:
```python
# create client and a new collection
db = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = db.get_or_create_collection("library_chat", embedding_function=LlamaCppEmbeddingFunction(model_path=temp_path, n_ctx=512, n_threads=n_cpu_cores, n_gpu_layers=32))
documents = SimpleDirectoryReader("./text/").load_data()
vector_store = ChromaVectorStore(chroma_collection=chroma_collection) # Create a vector store that uses the ChromaDB collection
storage_context = StorageContext.from_defaults(vector_store=vector_store) # Generates a storage context with default settings
index = VectorStoreIndex.from_documents(
documents, storage_context=storage_context # Embeds the documents using the LLM model
)
```
But I don't think it is currently using the embedding function I specified for it
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by locating the collections class's embedding-function list and checking how a supplied embedding_function is invoked. Review the Python embedding integration points, then verify the requested llama-cpp-python function is used for document and query embeddings. Done means a local Llama model can be selected through the collection API and the behavior is covered by tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- ai, database
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100