From 24fa515e9a9d4bd7c795d81aa7a8d0eef75e457e Mon Sep 17 00:00:00 2001 From: Bhanu Angam Date: Thu, 28 Sep 2023 12:15:38 +0200 Subject: [PATCH] Added use_auth_token to TextEmbedding class --- src/vecs/adapter/text.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/vecs/adapter/text.py b/src/vecs/adapter/text.py index 8e5c563..e2c6069 100644 --- a/src/vecs/adapter/text.py +++ b/src/vecs/adapter/text.py @@ -35,13 +35,14 @@ class TextEmbedding(AdapterStep): embeddings using a specified sentence transformers model. """ - def __init__(self, *, model: TextEmbeddingModel, batch_size: int = 8): + def __init__(self, *, model: TextEmbeddingModel, batch_size: int = 8, use_auth_token: str = None): """ Initializes the TextEmbedding adapter with a sentence transformers model. Args: model (TextEmbeddingModel): The sentence transformers model to use for embeddings. batch_size (int): The number of records to encode simultaneously. + use_auth_token (str): The HuggingFace Hub auth token to use for private models. Raises: MissingDependency: If the sentence_transformers library is not installed. @@ -53,7 +54,7 @@ def __init__(self, *, model: TextEmbeddingModel, batch_size: int = 8): "Missing feature vecs[text_embedding]. Hint: `pip install 'vecs[text_embedding]'`" ) - self.model = ST(model) + self.model = ST(model, use_auth_token=use_auth_token) self._exported_dimension = self.model.get_sentence_embedding_dimension() self.batch_size = batch_size