diff --git a/docker/compose.yaml b/docker/compose.yaml index f536f0c..1a5ddb5 100644 --- a/docker/compose.yaml +++ b/docker/compose.yaml @@ -31,6 +31,7 @@ services: image: "kemingy/spladepp" environment: - MOSEC_TIMEOUT=10000 + - HF_TOKEN=hf_xxxxxxxxx # add your huggingface token here (required by https://huggingface.co/prithivida/Splade_PP_en_v1) ports: - "8083:8000" diff --git a/pyproject.toml b/pyproject.toml index af60d5b..675eff3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "httpx~=0.27", "cohere~=4.45", "prometheus-client~=0.20", + "numpy~=1.26", ] [project.optional-dependencies] dev = [ diff --git a/qtext/metrics.py b/qtext/metrics.py index df06879..ea8aff7 100644 --- a/qtext/metrics.py +++ b/qtext/metrics.py @@ -7,6 +7,9 @@ doc_counter = Counter("add_doc", "Added documents", labelnames=labels) embedding_histogram = Histogram("embedding_latency_seconds", "Embedding cost time") sparse_histogram = Histogram("sparse_latency_seconds", "Sparse embedding cost time") +add_doc_histogram = Histogram( + "add_doc_latency_seconds", "Add doc cost time", labelnames=labels +) text_search_histogram = Histogram( "full_text_search_latency_seconds", "Full text search cost time", diff --git a/qtext/pg_client.py b/qtext/pg_client.py index 33ffa38..144e965 100644 --- a/qtext/pg_client.py +++ b/qtext/pg_client.py @@ -17,6 +17,7 @@ sparse_search_histogram, text_search_histogram, vector_search_histogram, + add_doc_histogram, ) from qtext.schema import DefaultTable, Querier from qtext.spec import AddNamespaceRequest, QueryDocRequest, SparseEmbedding @@ -141,6 +142,7 @@ def add_doc(self, req): if primary_id is None or getattr(req, primary_id, None) is None: attributes.remove(primary_id) placeholders = [getattr(req, key) for key in attributes] + start_time = perf_counter() self.conn.execute( sql.SQL( "INSERT INTO {table} ({fields}) VALUES ({placeholders})" @@ -154,6 +156,7 @@ def add_doc(self, req): placeholders, ) self.conn.commit() + add_doc_histogram.labels(req.namespace).observe(perf_counter() - start_time) doc_counter.labels(req.namespace).inc() except psycopg.errors.Error as err: logger.info("pg client add doc error", exc_info=err)