mirror of
https://github.com/neuml/txtai.git
synced 2026-10-12 02:19:52 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e981d2d8fc | ||
|
|
55fadb4c56 | ||
|
|
fe5be99bfa | ||
|
|
243aaa191d | ||
|
|
93a767f004 | ||
|
|
62099b9a34 | ||
|
|
c786f1b5ec | ||
|
|
738df97d5f | ||
|
|
4605ee001f | ||
|
|
46516fa4e6 | ||
|
|
fe1fbbc933 | ||
|
|
084a1eca4d | ||
|
|
ba9d52310a | ||
|
|
ca9563daac | ||
|
|
0e7202caee | ||
|
|
c1f1686889 | ||
|
|
04f5bf1998 | ||
|
|
c8e7d813af | ||
|
|
edf61d1d4a | ||
|
|
7e92f5a7d7 | ||
|
|
b2ff9bd85c | ||
|
|
9c718f7876 | ||
|
|
c486f4706e | ||
|
|
341184ba46 | ||
|
|
647d80f95e | ||
|
|
408c25d201 | ||
|
|
b0f29a11d1 | ||
|
|
49a64943b8 | ||
|
|
af33f2e024 | ||
|
|
f908b5d88d | ||
|
|
2e8493dfca | ||
|
|
d7e3e29ac3 | ||
|
|
4e2efa2cbc | ||
|
|
ffcbc0b815 | ||
|
|
f14c26c4c7 | ||
|
|
346bb9d282 | ||
|
|
e3e0413747 | ||
|
|
32655393e0 | ||
|
|
9118875980 | ||
|
|
d9b4476187 | ||
|
|
ff5aec468a | ||
|
|
3446a4ef13 | ||
|
|
84aa1c81c7 | ||
|
|
6bc76132d0 | ||
|
|
893d0424be | ||
|
|
98c965cc3b | ||
|
|
00622f2cf1 | ||
|
|
55d52fda92 | ||
|
|
20c9b8e0be | ||
|
|
3f33fb2918 | ||
|
|
3c7a548923 | ||
|
|
c78fbd976c | ||
|
|
2d535ff215 | ||
|
|
6765330f4a | ||
|
|
1d5041b6e7 | ||
|
|
d467878064 | ||
|
|
3bc07b3f09 | ||
|
|
8ef454d0bf | ||
|
|
47d3e65756 | ||
|
|
81f908874b | ||
|
|
bed91a0b9f | ||
|
|
7df50abb61 | ||
|
|
80b8ceff4a | ||
|
|
de9e09e9ae | ||
|
|
69ede9ec3e | ||
|
|
f8689ed903 | ||
|
|
0abf3763b4 | ||
|
|
11b2cfe298 | ||
|
|
d0cc2bed9c | ||
|
|
a5ea35d7a7 | ||
|
|
6a15ed8490 | ||
|
|
76147369fb | ||
|
|
2e89af0845 | ||
|
|
d9d6a1cbb1 | ||
|
|
f87574de4c | ||
|
|
e313248ff2 | ||
|
|
5cba72c88a | ||
|
|
45d334f4ce | ||
|
|
95d4c1c881 | ||
|
|
ef28da8735 | ||
|
|
f1304d8403 | ||
|
|
97ee7e073a | ||
|
|
1d3f8e9ebf | ||
|
|
f3ceece76a | ||
|
|
6737ee4e01 | ||
|
|
e30e8edbcc | ||
|
|
483dbd3b3b | ||
|
|
c09a6f24f1 | ||
|
|
6e7f0d260e | ||
|
|
793ecc8dc2 | ||
|
|
37500887ae | ||
|
|
4e2eb242b6 | ||
|
|
83c75ea28a | ||
|
|
f47d389735 | ||
|
|
4bb2141954 | ||
|
|
5d8c999ae0 | ||
|
|
48a34e21c3 | ||
|
|
262b03027e | ||
|
|
36c375b86e | ||
|
|
b715812e68 | ||
|
|
dcdded892d | ||
|
|
e7894cdb00 | ||
|
|
9f9b32e070 | ||
|
|
00167b9f08 | ||
|
|
3c07766199 | ||
|
|
7cb6bd09db | ||
|
|
4b7d70e9e8 | ||
|
|
2ed4e46ebf | ||
|
|
47d5f56768 | ||
|
|
d5db643540 | ||
|
|
9f7b4ddd42 | ||
|
|
19ff237c37 | ||
|
|
ad77514b32 | ||
|
|
144a957d42 | ||
|
|
2e8a3d7892 | ||
|
|
72f9ca7605 | ||
|
|
ebe13ae304 | ||
|
|
5d0e995f7f | ||
|
|
7965a7f42e | ||
|
|
2bae6e4d43 | ||
|
|
767272f2d5 | ||
|
|
21460347dd | ||
|
|
f3945f9120 | ||
|
|
3a7b597d78 | ||
|
|
20e2dc41e1 | ||
|
|
eafe574d92 | ||
|
|
7e98746d78 | ||
|
|
6ebfbbb716 | ||
|
|
8fef6ac454 | ||
|
|
0eafc0580f | ||
|
|
5ade7881c8 | ||
|
|
948dc42d19 | ||
|
|
2e8af221f4 | ||
|
|
396b3b8065 | ||
|
|
80420ee10a | ||
|
|
1d45e88569 | ||
|
|
007d61077a | ||
|
|
84fcd81039 | ||
|
|
cfe2227f30 | ||
|
|
82b5320da3 | ||
|
|
18b93be702 | ||
|
|
027c07deb9 | ||
|
|
876cd07cf7 | ||
|
|
546e7777f1 | ||
|
|
27542ca64b | ||
|
|
f0ee7d2b2d | ||
|
|
bba7fd589e | ||
|
|
6349ca2eac | ||
|
|
ba1693480e | ||
|
|
e90c67ef06 | ||
|
|
57991cb164 | ||
|
|
c2b264d297 | ||
|
|
406e417be0 | ||
|
|
954b17af4d | ||
|
|
2a139449dc | ||
|
|
a10667a1c2 | ||
|
|
34e6efcdd6 | ||
|
|
9b7a96b4d1 | ||
|
|
b94f770301 | ||
|
|
adbcccda46 | ||
|
|
20f818f72c | ||
|
|
ee0416e87b | ||
|
|
00ec81a10e | ||
|
|
059ecec2ee | ||
|
|
21c2c38b98 | ||
|
|
47c05994bf | ||
|
|
8f3735101b | ||
|
|
438ade31b6 | ||
|
|
3ea197468d | ||
|
|
99ebbdfea2 | ||
|
|
5570dbc7b6 | ||
|
|
01b456900a | ||
|
|
a0857e980a | ||
|
|
84be3476d8 | ||
|
|
f35d932979 | ||
|
|
a1c9b94d00 | ||
|
|
11ebacdb23 | ||
|
|
5986fa6082 | ||
|
|
893a333887 | ||
|
|
e54255f4e4 | ||
|
|
f89835bc77 | ||
|
|
ea531b3420 | ||
|
|
756917f14b | ||
|
|
b1f4845284 | ||
|
|
d2b1d3634f | ||
|
|
193facd144 | ||
|
|
436662e32b | ||
|
|
54644a0526 | ||
|
|
817a2b142a | ||
|
|
987cff001f | ||
|
|
8383e6ce8d | ||
|
|
0be1231d8d | ||
|
|
79e277fb7d | ||
|
|
f4472312ba | ||
|
|
90f7cc2078 | ||
|
|
ce11b08ab1 | ||
|
|
75b4ce5c7f | ||
|
|
2ef401bb78 | ||
|
|
8c43c70687 | ||
|
|
48e22b593c | ||
|
|
3beea7f4bf | ||
|
|
626cc8fbe0 | ||
|
|
21b9049c86 | ||
|
|
abf504f70e | ||
|
|
d9e9a9e684 | ||
|
|
bd199158b6 | ||
|
|
390cc08b11 | ||
|
|
feac4c8306 | ||
|
|
d7f974b3d7 | ||
|
|
d8fa848f73 | ||
|
|
bbd23bb12c | ||
|
|
8362110c64 | ||
|
|
83e1835b6d | ||
|
|
cf2368269f | ||
|
|
a910ddbbea | ||
|
|
7b49504dca | ||
|
|
86af7ec344 | ||
|
|
50ca7b2fde | ||
|
|
7c8d001492 | ||
|
|
dea0a65fd1 | ||
|
|
a1224c06f1 | ||
|
|
162b65d3d9 | ||
|
|
363f8ca7cf | ||
|
|
1d084941f4 | ||
|
|
597ba8133d | ||
|
|
6bbb25872a | ||
|
|
f22bf5bbb5 | ||
|
|
83871c6fe3 | ||
|
|
2c6ec162d0 | ||
|
|
d59c88424d | ||
|
|
95c7b8df1e | ||
|
|
2389b5946c | ||
|
|
b246b642c7 | ||
|
|
46467a3453 | ||
|
|
fcfd340ff1 | ||
|
|
77965f44dd | ||
|
|
7435297bac | ||
|
|
3bb34eae2e | ||
|
|
79a28641b2 | ||
|
|
7c77c134e2 | ||
|
|
d6f696d795 | ||
|
|
d8ae6776f8 | ||
|
|
e18e3d1245 | ||
|
|
ee0d1a0214 | ||
|
|
5ddedf07cd | ||
|
|
590c71f593 | ||
|
|
c33188a0f8 | ||
|
|
638a2e12b1 | ||
|
|
03da5c96fb | ||
|
|
99086302ce | ||
|
|
2f11f93743 | ||
|
|
783d9cf680 | ||
|
|
281babeff7 | ||
|
|
a71732198a | ||
|
|
14cbe51e04 | ||
|
|
5d15f7fa66 | ||
|
|
ac2dc6a30b | ||
|
|
7791174c1b | ||
|
|
b21ba7065e | ||
|
|
10a6b23304 | ||
|
|
576f0ddcd7 | ||
|
|
9b8a606a16 | ||
|
|
01344b021c | ||
|
|
dd00608b9a | ||
|
|
702e0b880f | ||
|
|
14ac0e4cc5 | ||
|
|
1f28e5401d | ||
|
|
a84f0a388c | ||
|
|
510f9c18ed | ||
|
|
ae67f63359 | ||
|
|
0e16a22051 | ||
|
|
bfd0aa7ff4 | ||
|
|
ba2ea2db4e | ||
|
|
1daac761aa | ||
|
|
e9e65393a0 | ||
|
|
24ca4e36ca | ||
|
|
8208780b6a | ||
|
|
c3735dd5d4 | ||
|
|
6544d123f6 | ||
|
|
aa0a8cfc32 | ||
|
|
fc6f93828b | ||
|
|
796c5bb0a8 | ||
|
|
7613acd6e5 | ||
|
|
9db11953c5 |
@@ -10,7 +10,7 @@ jobs:
|
||||
matrix:
|
||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
||||
|
||||
timeout-minutes: 60
|
||||
timeout-minutes: 75
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v6
|
||||
@@ -18,7 +18,7 @@ jobs:
|
||||
- name: Install Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.10"
|
||||
python-version: "3.11"
|
||||
|
||||
- name: Install Java
|
||||
uses: actions/setup-java@v5
|
||||
@@ -40,6 +40,7 @@ jobs:
|
||||
echo "TIKA_STARTUP_SLEEP=30" >> $GITHUB_ENV
|
||||
echo "TIKA_STARTUP_MAX_RETRY=10" >> $GITHUB_ENV
|
||||
brew install portaudio
|
||||
pip install torch==2.12.1 torchvision==0.27.1
|
||||
if: matrix.os == 'macos-latest'
|
||||
|
||||
- name: Install dependencies - Windows
|
||||
@@ -60,12 +61,16 @@ jobs:
|
||||
HF_HUB_ETAG_TIMEOUT: 100
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 100
|
||||
HF_XET_CHUNK_CACHE_SIZE_BYTES: 0
|
||||
DOCLING_DEVICE: "cpu"
|
||||
DOCLING_INFERENCE_COMPILE_TORCH_MODELS: false
|
||||
|
||||
- uses: pre-commit/[email protected]
|
||||
if: matrix.os == 'ubuntu-latest'
|
||||
|
||||
- name: Test Coverage
|
||||
run: coveralls --service=github
|
||||
run: |
|
||||
coverage report --fail-under=100
|
||||
coveralls --service=github
|
||||
if: matrix.os == 'ubuntu-latest'
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
@@ -13,7 +13,7 @@ jobs:
|
||||
- name: Install Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.10"
|
||||
python-version: "3.11"
|
||||
|
||||
- name: Install txtai
|
||||
run: |
|
||||
|
||||
@@ -13,7 +13,7 @@ jobs:
|
||||
- name: Install Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.10"
|
||||
python-version: "3.11"
|
||||
|
||||
- name: Minimal install
|
||||
run: |
|
||||
|
||||
@@ -12,6 +12,8 @@ max-locals=40
|
||||
max-returns=10
|
||||
max-attributes=20
|
||||
min-public-methods=0
|
||||
max-public-methods=80
|
||||
|
||||
[FORMAT]
|
||||
max-line-length=150
|
||||
max-module-lines=1500
|
||||
|
||||
@@ -49,7 +49,7 @@ Summary of txtai features:
|
||||
- 🔋 Batteries included with defaults to get up and running fast
|
||||
- ☁️ Run local or scale out with container orchestration
|
||||
|
||||
txtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
|
||||
txtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
|
||||
|
||||
> [!NOTE]
|
||||
>
|
||||
@@ -190,7 +190,7 @@ The easiest way to install is via pip and PyPI
|
||||
pip install txtai
|
||||
```
|
||||
|
||||
Python 3.10+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
|
||||
Python 3.11+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
|
||||
|
||||
See the detailed [install instructions](https://neuml.github.io/txtai/install) for more information covering [optional dependencies](https://neuml.github.io/txtai/install/#optional-dependencies), [environment specific prerequisites](https://neuml.github.io/txtai/install/#environment-specific-prerequisites), [installing from source](https://neuml.github.io/txtai/install/#install-from-source), [conda support](https://neuml.github.io/txtai/install/#conda), [lightweight minimal installation](https://neuml.github.io/txtai/install/#minimal-install) and how to [run with containers](https://neuml.github.io/txtai/cloud).
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# Set base image
|
||||
ARG BASE_IMAGE=python:3.10-slim
|
||||
ARG BASE_IMAGE=python:3.11-slim
|
||||
FROM $BASE_IMAGE
|
||||
|
||||
# Install GPU-enabled version of PyTorch if set
|
||||
@@ -8,7 +8,7 @@ ARG GPU
|
||||
# Target CPU architecture
|
||||
ARG TARGETARCH
|
||||
|
||||
# Set Python version (i.e. 3, 3.10)
|
||||
# Set Python version (i.e. 3, 3.11)
|
||||
ARG PYTHON_VERSION=3
|
||||
|
||||
# List of txtai components to install
|
||||
@@ -27,7 +27,7 @@ RUN \
|
||||
# Install txtai project and dependencies
|
||||
ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python && \
|
||||
python -m pip install --no-cache-dir -U pip wheel setuptools && \
|
||||
if [ -z ${GPU} ]; then pip install --no-cache-dir torch==2.12.0+cpu torchvision==0.27.0+cpu -f https://download.pytorch.org/whl/torch -f https://download.pytorch.org/whl/torchvision; fi && \
|
||||
if [ -z ${GPU} ]; then pip install --no-cache-dir torch==2.14.1+cpu torchvision==0.29.1+cpu -f https://download.pytorch.org/whl/torch -f https://download.pytorch.org/whl/torchvision; fi && \
|
||||
python -m pip install --no-cache-dir txtai${COMPONENTS} && \
|
||||
python -c "import sys, importlib.util as util; 1 if util.find_spec('nltk') else sys.exit(); import nltk; nltk.download(['punkt', 'punkt_tab', 'averaged_perceptron_tagger_eng'])" && \
|
||||
\
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
# Set base image
|
||||
ARG BASE_IMAGE=python:3.10-slim
|
||||
ARG BASE_IMAGE=python:3.11-slim
|
||||
FROM $BASE_IMAGE
|
||||
|
||||
# Target CPU architecture
|
||||
ARG TARGETARCH
|
||||
|
||||
# Set Python version (i.e. 3, 3.10)
|
||||
# Set Python version (i.e. 3, 3.11)
|
||||
ARG PYTHON_VERSION=3
|
||||
|
||||
# Locale environment variables
|
||||
|
||||
+43
-2
@@ -6,9 +6,9 @@ The default implementation of an API service runs via HTTP and is fully open. If
|
||||
|
||||
The default API service command starts a Uvicorn server as a HTTP service on port 8000. To run a HTTPS service, consider the following options.
|
||||
|
||||
- [TLS Proxy Server](https://fastapi.tiangolo.com/deployment/https/). *Recommended choice*. With this configuration, the txtai API service runs as a HTTP service only accessible on the localhost/local network. The proxy server handles all encryption and redirects requests to local services. See this [example configuration](https://www.uvicorn.org/deployment/#running-behind-nginx) for more.
|
||||
- [TLS Proxy Server](https://fastapi.tiangolo.com/deployment/https/). *Recommended choice*. With this configuration, the txtai API service runs as a HTTP service only accessible on the localhost/local network. The proxy server handles all encryption and redirects requests to local services. See this [example configuration](https://www.uvicorn.dev/deployment/#running-behind-nginx) for more.
|
||||
|
||||
- [Uvicorn SSL Certificate](https://www.uvicorn.org/deployment/). Another option is setting the SSL certificate on the Uvicorn service. This works in simple situations but gets complex when hosting multiple txtai or other related services.
|
||||
- [Uvicorn SSL Certificate](https://www.uvicorn.dev/deployment/). Another option is setting the SSL certificate on the Uvicorn service. This works in simple situations but gets complex when hosting multiple txtai or other related services.
|
||||
|
||||
## Authorization
|
||||
|
||||
@@ -31,3 +31,44 @@ curl \
|
||||
It's important to note that HTTPS **must** be enabled using one of the methods mentioned above. Otherwise, tokens will be exchanged as clear text.
|
||||
|
||||
Authentication and Authorization can be fully customized. See the [dependencies](../customization#dependencies) section for more.
|
||||
|
||||
## Safe Open
|
||||
|
||||
Many components in txtai work with local files and have no restrictions on what files and urls can be accessed. The [textractor pipeline](../../pipeline/data/textractor) and [retrieve](../../workflow/task/retrieve) tasks both have a built-in option called `safeopen` which is enabled by default when running the API and applications. It enforces the following rules:
|
||||
|
||||
- Local files must be in the safeopen directory (defaults to temp dir)
|
||||
- URLs must be public URLs (also has checks for redirects and DNS rebind attempts)
|
||||
|
||||
The following workflow only allows accessing files stored in a specific directory along with public urls.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
textract:
|
||||
tasks:
|
||||
- action: textractor
|
||||
```
|
||||
|
||||
A number of of pipelines support reading local files. For these it's recommended to wrap those pipelines as workflows with a `retrieve` task. Note that pipelines can be created in-line and don't need to be exposed via the API.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
caption:
|
||||
tasks:
|
||||
- task: retrieve
|
||||
action: caption
|
||||
```
|
||||
|
||||
Additionally, internal pipelines can be created but not exposed via the API.
|
||||
|
||||
```yaml
|
||||
routes:
|
||||
caption: False
|
||||
|
||||
caption:
|
||||
|
||||
workflow:
|
||||
caption:
|
||||
tasks:
|
||||
- task: retrieve
|
||||
action: caption
|
||||
```
|
||||
|
||||
+2
-2
@@ -30,8 +30,8 @@ Examples build commands below.
|
||||
# Get Dockerfile
|
||||
wget https://raw.githubusercontent.com/neuml/txtai/master/docker/base/Dockerfile
|
||||
|
||||
# Build Ubuntu 22.04 image running Python 3.10
|
||||
docker build -t txtai --build-arg BASE_IMAGE=ubuntu:22.04 --build-arg PYTHON_VERSION=3.10 .
|
||||
# Build Ubuntu 22.04 image running Python 3.11
|
||||
docker build -t txtai --build-arg BASE_IMAGE=ubuntu:22.04 --build-arg PYTHON_VERSION=3.11 .
|
||||
|
||||
# Build image with GPU support
|
||||
docker build -t txtai --build-arg GPU=1 .
|
||||
|
||||
@@ -4,7 +4,7 @@ Approximate Nearest Neighbor (ANN) index configuration for storing vector embedd
|
||||
|
||||
## backend
|
||||
```yaml
|
||||
backend: faiss|hnsw|annoy|ggml|numpy|torch|turbovec|pgvector|sqlite|custom
|
||||
backend: faiss|hnsw|annoy|ggml|milvus|numpy|rabitq|torch|turbovec|zvec|pgvector|sqlite|custom
|
||||
```
|
||||
|
||||
Sets the ANN backend. Defaults to `faiss`. Additional backends are available via the [ann](../../../install/#ann) extras package. Set custom backends via setting this parameter to the fully resolvable class string.
|
||||
@@ -74,9 +74,14 @@ The [GGML](https://github.com/ggml-org/ggml) backend is a k-nearest neighbors ba
|
||||
|
||||
[See this](https://github.com/ggml-org/ggml/blob/master/include/ggml.h#L379) for a list of quantization types.
|
||||
|
||||
### numpy
|
||||
### milvus
|
||||
|
||||
The NumPy backend is a k-nearest neighbors backend. It's designed for simplicity and works well with smaller datasets that fit into memory.
|
||||
```yaml
|
||||
milvus:
|
||||
m: number of HNSW links per element (int) - defaults to 50
|
||||
```
|
||||
|
||||
### numpy
|
||||
|
||||
```yaml
|
||||
numpy:
|
||||
@@ -84,12 +89,30 @@ numpy:
|
||||
defaults to NumPy array storage
|
||||
```
|
||||
|
||||
The NumPy backend is a k-nearest neighbors backend. It's designed for simplicity and works well with smaller datasets that fit into memory.
|
||||
|
||||
### rabitq
|
||||
|
||||
```yaml
|
||||
rabitq:
|
||||
mode: index mode (ivf or hnsw) - defaults to "ivf"
|
||||
nbits: number of quantization bits per dimension (int) - defaults to 1,
|
||||
supports 1 to 9 in both modes, ivf mode also supports 32 which
|
||||
keeps raw vectors for reranking
|
||||
clusters: number of IVF clusters (int) - defaults to
|
||||
max(1, min(round(4 * sqrt(embeddings count)), embeddings count))
|
||||
nprobe: search probe setting for ivf mode (int) - defaults to
|
||||
max(1, round(num_clusters/16))
|
||||
m: M param for hnsw mode (int) - defaults to 16
|
||||
efconstruction: ef_construction param for hnsw mode (int) - defaults to 200
|
||||
efsearch: ef search param for hnsw mode (int) - defaults to None and not set
|
||||
randomseed: random-seed param for hnsw mode (int) - defaults to 100
|
||||
```
|
||||
|
||||
The [rabitq](https://github.com/VectorDB-NTU/RaBitQ-Library) backend is a quantized index powered by the [RaBitQ algorithm](https://doi.org/10.1145/3725413). It supports ivf and hnsw search modes. Vectors are stored with 1-bit quantization by default. More bits increase accuracy at the cost of memory. The index is stored as a single file and no vectors are kept outside of it. With nbits set to 32, the vectors are stored in the index for reranking. The ivf mode supports upserts and deletes. Deleted rows keep their storage. Rows added after the first index are assigned to the clusters trained when the index was built, so recall can drop if the new data differs from the original data. Rebuild the index to train new clusters. Note that hnsw indexes can not be modified after creation, upserts/deletes and other modifications are not supported.
|
||||
|
||||
### torch
|
||||
|
||||
The Torch backend is a k-nearest neighbors backend like NumPy. It supports GPU-enabled operations. It also has support for quantization which enables fitting larger arrays into GPU memory.
|
||||
|
||||
When quantization is enabled, vectors are _always_ stored in safetensors. _Note that macOS support for quantization is limited._
|
||||
|
||||
```yaml
|
||||
torch:
|
||||
safetensors: stores vectors using the safetensors format - defaults
|
||||
@@ -99,13 +122,29 @@ torch:
|
||||
blocksize: quantization block size parameter
|
||||
```
|
||||
|
||||
### turbovec
|
||||
The Torch backend is a k-nearest neighbors backend like NumPy. It supports GPU-enabled operations. It also has support for quantization which enables fitting larger arrays into GPU memory.
|
||||
|
||||
The [turbovec](https://github.com/RyanCodrai/turbovec) backend is a k-nearest neighbors backend powered by the [TurboQuant algorithm](https://arxiv.org/abs/2504.19874).
|
||||
When quantization is enabled, vectors are _always_ stored in safetensors. _Note that macOS support for quantization is limited._
|
||||
|
||||
### turbovec
|
||||
|
||||
```yaml
|
||||
turbovec:
|
||||
bitwidth: number of bits to store each vector dimension as. Supports 2, 3 or 4.
|
||||
```
|
||||
|
||||
The [turbovec](https://github.com/RyanCodrai/turbovec) backend is a k-nearest neighbors backend powered by the [TurboQuant algorithm](https://arxiv.org/abs/2504.19874).
|
||||
|
||||
### zvec
|
||||
|
||||
```yaml
|
||||
zvec:
|
||||
efconstruction: ef_construction param for HnswIndexParam (int) - defaults to 200
|
||||
m: number of HNSW links per element (int) - defaults to 50
|
||||
efsearch: ef search param for HnswQueryParam (int) - defaults to 300
|
||||
```
|
||||
|
||||
The [zvec](https://github.com/alibaba/zvec) backend is an embedded, path-based vector index.
|
||||
|
||||
### pgvector
|
||||
```yaml
|
||||
|
||||
@@ -36,7 +36,7 @@ batch: Sets the transform batch size
|
||||
encodebatch: Sets the encode batch size
|
||||
vectors: additional model init args
|
||||
encodeargs: additional encode() args
|
||||
backend: ivfsparse|pgsparse
|
||||
backend: ivfsparse|pgsparse|zvecsparse
|
||||
```
|
||||
|
||||
Sparse vector scoring options. The sparse scoring instance combines a sparse vector model with a sparse approximate nearest neighbor index (ANN). This method supports both vector normalization and score normalization.
|
||||
@@ -54,7 +54,7 @@ Score normalization scales the output between 0 and 1. This setting supports:
|
||||
```yaml
|
||||
ivfsparse:
|
||||
sample: percent of data to use for model training (0.0 - 1.0)
|
||||
nfeatures: top n features to use for model training (int)
|
||||
nfeatures: top n features to use for model training, defaults to all features (int)
|
||||
nlist: desired number of clusters (int)
|
||||
nprobe: search probe setting (int)
|
||||
minpoints: minimum number of points for a cluster (int)
|
||||
@@ -66,6 +66,16 @@ Inverted file (IVF) index with flat vector file storage and sparse array support
|
||||
|
||||
Sparse ANN backed by Postgres. Supports same options as the [pgvector](../ann/#pgvector) ANN.
|
||||
|
||||
#### zvecsparse
|
||||
```yaml
|
||||
zvecsparse:
|
||||
efconstruction: ef_construction param for HnswIndexParam (int) - defaults to 200
|
||||
m: number of HNSW links per element (int) - defaults to 50
|
||||
efsearch: ef search param for HnswQueryParam (int) - defaults to 300
|
||||
```
|
||||
|
||||
Sparse ANN backed by [zvec](https://github.com/alibaba/zvec). Stores sparse vectors in an embedded, path-based vector index. Requires the [ann](../../../install/#ann) extras package.
|
||||
|
||||
## terms
|
||||
```yaml
|
||||
terms: boolean|dict
|
||||
@@ -81,6 +91,15 @@ Supports a `dict` with the parameters `cachelimit` and `cutoff`.
|
||||
|
||||
When `terms` is set to `True`, default parameters are used for the `cachelimit` and `cutoff`. Normally, these defaults are sufficient.
|
||||
|
||||
### Wildcard queries
|
||||
|
||||
Keyword queries containing `*` expand matching terms using SQL `LIKE`. Within these
|
||||
expressions, `%` matches any sequence and `_` matches one character. To match a literal
|
||||
percent sign, underscore or backslash, use `\%`, `\_` or `\\`, respectively.
|
||||
For example, `report\_*` matches `report_2026`, while `report_*` also matches `reportX2026`.
|
||||
Use a tokenizer that preserves these characters, such as `tokenizer: {whitespace: true}`,
|
||||
or pass pre-tokenized query terms. Queries without `*` retain exact-term matching.
|
||||
|
||||
## normalize
|
||||
```yaml
|
||||
normalize: boolean|str|dict
|
||||
|
||||
@@ -25,8 +25,7 @@ Embeddings method to use. If the method is not provided, it is inferred using th
|
||||
Builds embeddings using a transformers model. While this can be any transformers model, it works best with
|
||||
[models trained](https://huggingface.co/models?pipeline_tag=sentence-similarity) to build embeddings.
|
||||
|
||||
`mean`, `cls`, `last` and `late` pooling are supported and automatically inferred from the model. The pooling method can be overwritten by changing the method
|
||||
from `transformers` to `meanpooling`, `clspooling`, `lastpooling` or `latepooling` respectively.
|
||||
`mean`, `cls`, `last`, `late` and `max` pooling are supported and automatically inferred from the model. The pooling method can be overwritten by changing the method from `transformers` to `meanpooling`, `clspooling`, `lastpooling`, `latepooling` or `maxpooling` respectively.
|
||||
|
||||
Setting `maxlength` to `True` enables truncating inputs to the `max_seq_length`. Setting `maxlength` to an integer will truncate inputs to that value. When omitted (default), the `maxlength` will be set to either the model or tokenizer maxlength.
|
||||
|
||||
@@ -68,7 +67,9 @@ Embeddings are created via an external model or API. Requires setting the [trans
|
||||
transform: function
|
||||
```
|
||||
|
||||
When method is `external`, this function transforms input content into embeddings. The input to this function is a list of data. This method must return either a numpy array or list of numpy arrays.
|
||||
When method is `external`, this function transforms input content into embeddings. The input to this function is a list of data. This method must return either a numpy array, list of numpy arrays or list of float arrays.
|
||||
|
||||
Note that this method is deprecated and will raise an error with the default settings. Adding a custom `Vectors` implementation is a cleaner approach. The `ALLOW_RESOLVE_TRANSFORM` environment variable must be set to enable. Only enable for trusted and/or reviewed code.
|
||||
|
||||
## gpu
|
||||
```yaml
|
||||
@@ -107,7 +108,7 @@ quantize: int|boolean
|
||||
|
||||
Enables scalar vector quantization at the specified precision. Supports 1-bit through 8-bit quantization. Scalar quantization transforms continuous floating point values to discrete unsigned integers. The `faiss`, `pgvector`, `numpy` and `torch` ANN backends support storing these vectors.
|
||||
|
||||
This parameter supports booleans for backwards compatability. When set to true/false, this flag sets [faiss.quantize](../ann/#faiss).
|
||||
This parameter supports booleans for backwards compatibility. When set to true/false, this flag sets [faiss.quantize](../ann/#faiss).
|
||||
|
||||
In addition to vector-level quantization, some ANN backends have the ability to quantize vectors at the storage layer. See the [ANN](../ann) configuration options for more.
|
||||
|
||||
@@ -143,6 +144,29 @@ vectors: dict
|
||||
|
||||
Passes these additional parameters to the underlying vector model.
|
||||
|
||||
### center
|
||||
```yaml
|
||||
center: boolean | dict
|
||||
```
|
||||
|
||||
Centers normalized late-interaction token vectors and normalizes them again before MUVERA or LEMUR encoding. Centering is applied to
|
||||
both query and document vectors. `true` uses batch scope and `false` disables centering. An explicit setting always takes precedence.
|
||||
When omitted, a LEMUR artifact with a stored collection mean uses collection scope. Other LEMUR artifacts and late-interaction models
|
||||
default to batch scope only when they load more than one `torch.nn.Linear` layer; models with zero or one loaded linear layer keep their
|
||||
previous output.
|
||||
|
||||
Dictionary configuration supports the following scopes:
|
||||
|
||||
```yaml
|
||||
center:
|
||||
scope: document | batch | collection
|
||||
```
|
||||
|
||||
`document` subtracts the mean of each document or query independently. `batch` subtracts the mean of all real token rows in the current
|
||||
model batch. Explicit `collection` scope requires a caller-provided, one-dimensional mean vector through either `mean` (an array-like
|
||||
value) or `path` (a safetensors file containing a `center.mean` tensor), but not both. LEMUR artifacts can provide this mean
|
||||
automatically. Zero-padding rows are never centered and remain zero.
|
||||
|
||||
### muvera
|
||||
```yaml
|
||||
muvera:
|
||||
@@ -152,6 +176,21 @@ muvera:
|
||||
```
|
||||
|
||||
Settings to control the size of MUVERA fixed dimensional outputs. Default is 20 * 2^5 * 16 = 10,240 dimensions.
|
||||
Set `muvera` to `false` to disable it or `true` to use the default settings. `false` only makes sense when the multi-vector output
|
||||
is consumed directly, for example through `PoolingFactory`. Embeddings vectors require `muvera` or `lemur`.
|
||||
|
||||
### lemur
|
||||
```yaml
|
||||
lemur:
|
||||
path: local artifact directory or Hugging Face Hub path
|
||||
```
|
||||
|
||||
Loads a trained [LEMUR](../../../pipeline/train/lemur) fixed dimensional encoder. LEMUR artifacts are corpus-specific and must be
|
||||
created with `LemurTrainer` before configuring an embeddings index. A bare string is treated as the artifact path. Query vectors are
|
||||
summed learned features and document vectors are ordinary least squares weights over the artifact's stored token sample. Each artifact
|
||||
contains `config.json` and `model.safetensors`; newly trained artifacts also store the collection token mean as `lemur.center`. Loading that mean automatically
|
||||
selects collection centering unless `center` is explicitly configured. Artifacts without `lemur.center` retain the previous default.
|
||||
The Safetensors file contains inference state only. See the trainer page for score filtering and Faiss exact/IVF search behavior.
|
||||
|
||||
### trust_remote_code
|
||||
```yaml
|
||||
|
||||
@@ -54,6 +54,8 @@ The number of candidates should be larger than the desired number of results whe
|
||||
- For a single query filter clause, the default is the query limit
|
||||
- With multiple filtering clauses, the default is 10x the query limit
|
||||
|
||||
When a query orders results by score, its `offset` adds to the query limit.
|
||||
|
||||
The index name is only applicable when [subindexes](../configuration/general/#indexes) are enabled. This specifies the index to use for the query.
|
||||
|
||||
Weights sets the hybrid score weights when an index has both a sparse and dense index.
|
||||
|
||||
+1
-1
@@ -51,7 +51,7 @@ Summary of txtai features:
|
||||
- 🔋 Batteries included with defaults to get up and running fast
|
||||
- ☁️ Run local or scale out with container orchestration
|
||||
|
||||
txtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
|
||||
txtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
|
||||
|
||||
!!! note
|
||||
|
||||
|
||||
+1
-1
@@ -9,7 +9,7 @@ The easiest way to install is via pip and PyPI
|
||||
pip install txtai
|
||||
```
|
||||
|
||||
Python 3.10+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
|
||||
Python 3.11+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
|
||||
|
||||
## Optional dependencies
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
# LemurTrainer
|
||||
|
||||

|
||||

|
||||
|
||||
Trains a LEMUR fixed dimensional encoder for a late interaction model and corpus. Training is a separate pipeline step; the saved
|
||||
artifact is then loaded through the embeddings `vectors.lemur` configuration.
|
||||
|
||||
`LemurTrainer` owns training-data construction, feature-model setup, epoch training and validation-based selection. The loaded LEMUR
|
||||
pooler remains responsible for query-feature and document-weight encoding.
|
||||
|
||||
## Example
|
||||
|
||||
```python
|
||||
from txtai.pipeline import LemurTrainer
|
||||
|
||||
corpus = [
|
||||
"First document",
|
||||
"Second document",
|
||||
"Third document",
|
||||
]
|
||||
|
||||
trainer = LemurTrainer()
|
||||
trainer(
|
||||
"colbert-ir/colbertv2.0",
|
||||
corpus,
|
||||
"lemur-model",
|
||||
epochs=100,
|
||||
validationsplit=0.1,
|
||||
)
|
||||
```
|
||||
|
||||
`epochs` is required so the training mode is explicit. Use `epochs=100` for the quality-oriented MLP setting. This can take hours on
|
||||
a CPU. Use `epochs=0` to select deterministic random ELM features as a lower-cost fallback.
|
||||
|
||||
The learn distribution defaults to `learncategory="query"`. Target documents are always encoded with the data encoder, so the
|
||||
default encodes selected corpus texts once as data and once as queries. Set `learncategory="data"` to reuse the data encodings, or
|
||||
pass a separate iterable of texts with `learn`.
|
||||
|
||||
When `vectors.center` is omitted, the trainer encodes the corpus without centering, computes one mean over all corpus token rows and
|
||||
uses that mean to center both data and learn vectors before fitting. The collection mean is stored in the LEMUR artifact. An explicit
|
||||
`vectors.center` setting keeps its configured behavior and does not store an automatic collection mean.
|
||||
|
||||
Set `corpussubsetsize` to a positive integer to sample that many raw corpus texts under `seed` before either encoding pass. The
|
||||
default is `None`, which uses every corpus text. `trainsubsetsize` and `learnsubsetsize` apply later, after token vectors have already
|
||||
been created. `corpussubsetsize` applies to `data`; a separate `learn` iterable remains caller-sized.
|
||||
|
||||
For trained MLP features, set `validationsplit` to a fraction greater than zero and less than one to retain the epoch with the lowest
|
||||
held-out loss. The default is `0.0`, which preserves training-loss selection. The selected one-based epoch, loss and metric are
|
||||
available as `selectedepoch`, `selectedloss` and `selectionmetric` on the fitted or reloaded encoder.
|
||||
|
||||
MLP training displays a `tqdm` progress bar on an interactive terminal, including percent complete and the current validation loss
|
||||
when `validationsplit` is enabled (or training loss otherwise). Progress output is disabled for non-interactive runs.
|
||||
|
||||
The artifact contains `config.json` and `model.safetensors`. It stores the inference feature model, output-normalization statistics,
|
||||
token sample needed to encode documents added after training and, by default, the collection token mean. The training-only output
|
||||
readout is not saved.
|
||||
|
||||
Load the artifact in an embeddings configuration.
|
||||
|
||||
```yaml
|
||||
embeddings:
|
||||
path: colbert-ir/colbertv2.0
|
||||
vectors:
|
||||
lemur:
|
||||
path: lemur-model
|
||||
```
|
||||
|
||||
## Search behavior
|
||||
|
||||
Loading a LEMUR artifact with a stored collection mean automatically centers both query and document token vectors with that mean.
|
||||
This keeps training and search on the same representation and makes fixed vectors independent of indexing or query batch size. An
|
||||
explicit `vectors.center` setting takes precedence. Older artifacts without a stored mean use the standard late-pooling fallback:
|
||||
batch centering for models with multiple linear layers and no centering for models with zero or one linear layer.
|
||||
|
||||
LEMUR approximates standardized MaxSim targets, so useful ranking scores can be negative. txtai's dense vector path L2-normalizes the
|
||||
fixed vectors and removes results with scores less than or equal to zero. This can change ordering and return fewer than the requested
|
||||
number of candidates compared with raw maximum inner-product search.
|
||||
|
||||
The Faiss backend uses exact `IDMap,Flat` search through 5,000 rows and switches to an IVF index above that threshold. In the measured
|
||||
scifact run, default IVF reduced LEMUR NDCG@10 by 43% relative to exact search, compared with 25% for MUVERA. For LEMUR corpora above
|
||||
5,000 rows, either pin `faiss.components` to an exact index or tune the IVF settings for the corpus. The exact configuration is:
|
||||
|
||||
```yaml
|
||||
embeddings:
|
||||
path: colbert-ir/colbertv2.0
|
||||
vectors:
|
||||
lemur:
|
||||
path: lemur-model
|
||||
faiss:
|
||||
components: IDMap,Flat
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
Python documentation for the pipeline.
|
||||
|
||||
### ::: txtai.pipeline.LemurTrainer.__call__
|
||||
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: console
|
||||
name:
|
||||
tasks:
|
||||
- task: console
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: export
|
||||
name:
|
||||
tasks:
|
||||
- task: export
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: file
|
||||
name:
|
||||
tasks:
|
||||
- task: file
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: image
|
||||
name:
|
||||
tasks:
|
||||
- task: image
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -35,12 +35,26 @@ Tasks can also be created with configuration as part of a workflow.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- action: summary
|
||||
name:
|
||||
tasks:
|
||||
- action: summary
|
||||
```
|
||||
|
||||
::: txtai.workflow.Task.__init__
|
||||
|
||||
## Multi-action task inputs
|
||||
|
||||
A multi-action `Task` requires reusable inputs. Passing a single-pass iterator, such as a generator or the output of a `StreamTask`, raises a `ValueError` before consuming inputs or running any actions. Automatically buffering an iterator could exhaust memory, so tasks do not convert these inputs to a list.
|
||||
|
||||
Iterators passed to a `Workflow` remain supported: the workflow collects bounded batches before passing them to its tasks. For example:
|
||||
|
||||
```python
|
||||
workflow = Workflow([Task([list, list])], batch=2)
|
||||
list(workflow(iter([1, 2, 3]))) # [(1, 1), (2, 2), (3, 3)]
|
||||
```
|
||||
|
||||
This batching occurs at the workflow input; it does not batch iterator outputs produced by tasks inside that workflow. Single-action tasks continue to accept iterator inputs.
|
||||
|
||||
## Multi-action task concurrency
|
||||
|
||||
The default processing mode is to run actions sequentially. Multiprocessing support is already built in at a number of levels. Any of the GPU models will maximize GPU utilization for example and even in CPU mode, concurrency is utilized. But there are still use cases for task action concurrency. For example, if the system has multiple GPUs, the task runs external sequential code, or the task has a large number of I/O tasks.
|
||||
@@ -72,7 +86,7 @@ workflow = Workflow([Task(lambda x: [y * 3 for y in x], unpack=False, column=0)]
|
||||
list(workflow([(2, 8)]))
|
||||
```
|
||||
|
||||
For the example input tuple of (2, 2), the workflow will only select the first element (2) and run the task against that element.
|
||||
For the example input tuple of (2, 8), the workflow will only select the first element (2) and run the task against that element.
|
||||
|
||||
```python
|
||||
workflow = Workflow([Task([lambda x: [y * 3 for y in x],
|
||||
|
||||
@@ -22,9 +22,10 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: retrieve
|
||||
directory: /tmp
|
||||
name:
|
||||
tasks:
|
||||
- task: retrieve
|
||||
directory: /tmp
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -22,9 +22,10 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: service
|
||||
url: https://service.url/action
|
||||
name:
|
||||
tasks:
|
||||
- task: service
|
||||
url: https://service.url/action
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: storage
|
||||
name:
|
||||
tasks:
|
||||
- task: storage
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -23,9 +23,10 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: template
|
||||
template: This is a {text} task
|
||||
name:
|
||||
tasks:
|
||||
- task: template
|
||||
template: This is a {text} task
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
|
||||
|
||||
```yaml
|
||||
workflow:
|
||||
tasks:
|
||||
- task: url
|
||||
name:
|
||||
tasks:
|
||||
- task: url
|
||||
```
|
||||
|
||||
## Methods
|
||||
|
||||
@@ -27,7 +27,7 @@
|
||||
"- 🔋 Batteries included with defaults to get up and running fast\n",
|
||||
"- ☁️ Run local or scale out with container orchestration\n",
|
||||
"\n",
|
||||
"txtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.\n",
|
||||
"txtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.\n",
|
||||
"\n",
|
||||
"> [NeuML](https://neuml.com) is the company behind txtai and we provide AI consulting services around our stack. [Schedule a meeting](https://cal.com/neuml/intro) or [send a message](mailto:[email protected]) to learn more.\n",
|
||||
">\n",
|
||||
|
||||
@@ -125,7 +125,7 @@
|
||||
"## Install\n",
|
||||
"The easiest way to install is via pip and PyPI\n",
|
||||
"pip install txtai\n",
|
||||
"Python 3.10+ is supported. Using a Python virtual environment is **recommended** .\n",
|
||||
"Python 3.11+ is supported. Using a Python virtual environment is **recommended** .\n",
|
||||
"See the detailed install instructions for more information covering optional dependencies, environment specific prerequisites, installing from source, conda support and how to run with containers.\n",
|
||||
"\n",
|
||||
"\n",
|
||||
@@ -525,8 +525,8 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"ANSWER: Python 3.10 and later versions (Python 3.10+) are supported.\n",
|
||||
"CITATION: [{'id': '24', 'text': 'Python 3.10+ is supported. Using a Python virtual environment is recommended.'}]\n"
|
||||
"ANSWER: Python 3.11 and later versions (Python 3.11+) are supported.\n",
|
||||
"CITATION: [{'id': '24', 'text': 'Python 3.11+ is supported. Using a Python virtual environment is recommended.'}]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
|
||||
@@ -100,12 +100,12 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"{'indexid': 0, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '**GitHub - neuml/txtai: 💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows**\\n\\n*💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows - neuml/txtai*\\n\\n\\n\\n**All-in-one embeddings database** \\ntxtai is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\\n\\nEmbeddings databases are a union of vector indexes (sparse and dense), graph networks and relational databases.\\n\\nThis foundation enables vector search and/or serves as a powerful knowledge source for large language model (LLM) applications.\\n\\nBuild autonomous agents, retrieval augmented generation (RAG) processes, multi-model workflows and more.\\n\\nSummary of txtai features:\\n\\n- 🔎 Vector search with SQL, object storage, topic modeling, graph analysis and multimodal indexing\\n- 📄 Create embeddings for text, documents, audio, images and video\\n- 💡 Pipelines powered by language models that run LLM prompts, question-answering, labeling, transcription, translation, summarization and more\\n- ↪️️ Workflows to join pipelines together and aggregate business logic. txtai processes can be simple microservices or multi-model workflows.\\n- 🤖 Agents that intelligently connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems\\n- ⚙️ Build with Python or YAML. API bindings available for [JavaScript](https://github.com/neuml/txtai.js) , [Java](https://github.com/neuml/txtai.java) , [Rust](https://github.com/neuml/txtai.rs) and [Go](https://github.com/neuml/txtai.go) .\\n- 🔋 Batteries included with defaults to get up and running fast\\n- ☁️ Run local or scale out with container orchestration\\ntxtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers) , [Sentence Transformers](https://github.'}\n",
|
||||
"{'indexid': 0, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '**GitHub - neuml/txtai: 💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows**\\n\\n*💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows - neuml/txtai*\\n\\n\\n\\n**All-in-one embeddings database** \\ntxtai is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\\n\\nEmbeddings databases are a union of vector indexes (sparse and dense), graph networks and relational databases.\\n\\nThis foundation enables vector search and/or serves as a powerful knowledge source for large language model (LLM) applications.\\n\\nBuild autonomous agents, retrieval augmented generation (RAG) processes, multi-model workflows and more.\\n\\nSummary of txtai features:\\n\\n- 🔎 Vector search with SQL, object storage, topic modeling, graph analysis and multimodal indexing\\n- 📄 Create embeddings for text, documents, audio, images and video\\n- 💡 Pipelines powered by language models that run LLM prompts, question-answering, labeling, transcription, translation, summarization and more\\n- ↪️️ Workflows to join pipelines together and aggregate business logic. txtai processes can be simple microservices or multi-model workflows.\\n- 🤖 Agents that intelligently connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems\\n- ⚙️ Build with Python or YAML. API bindings available for [JavaScript](https://github.com/neuml/txtai.js) , [Java](https://github.com/neuml/txtai.java) , [Rust](https://github.com/neuml/txtai.rs) and [Go](https://github.com/neuml/txtai.go) .\\n- 🔋 Batteries included with defaults to get up and running fast\\n- ☁️ Run local or scale out with container orchestration\\ntxtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers) , [Sentence Transformers](https://github.'}\n",
|
||||
"{'indexid': 1, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\\n\\n*Interested in an easy and secure way to run hosted txtai applications? Then join the* [txtai.cloud](https://txtai.cloud) *preview to learn more.* \\n\\n## Why txtai?\\nNew vector databases, LLM frameworks and everything in between are sprouting up daily. Why build with txtai?\\n\\n- Up and running in minutes with [pip](https://neuml.github.io/txtai/install/) or [Docker](https://neuml.github.io/txtai/cloud/) \\n```\\n# Get started in a couple lines\\nimport txtai\\n\\nembeddings = txtai.Embeddings()\\nembeddings.index([\"Correct\", \"Not what we hoped\"])\\nembeddings.search(\"positive\", 1)\\n#[(0, 0.29862046241760254)]\\n```\\n\\n- Built-in API makes it easy to develop applications using your programming language of choice\\n```\\n# app.yml\\nembeddings:\\n path: sentence-transformers/all-MiniLM-L6-v2\\n```\\n\\n```\\nCONFIG=app.yml uvicorn \"txtai.api:app\"\\ncurl -X GET \"http://localhost:8000/search?query=positive\"\\n```\\n\\n- Run local - no need to ship data off to disparate remote services\\n- Work with micromodels all the way up to large language models (LLMs)\\n- Low footprint - install additional dependencies and scale up when needed\\n- [Learn by example](https://neuml.github.io/txtai/examples) - notebooks cover all available functionality\\n\\n## Use Cases\\nThe following sections introduce common txtai use cases. A comprehensive set of over 60 [example notebooks and applications](https://neuml.github.io/txtai/examples) are also available.\\n\\n\\n### Semantic Search\\nBuild semantic/similarity/vector/neural search applications.'}\n",
|
||||
"{'indexid': 2, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'Traditional search systems use keywords to find data. Semantic search has an understanding of natural language and identifies results that have the same meaning, not necessarily the same keywords.\\n\\nGet started with the following examples.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Introducing txtai](https://github.com/neuml/txtai/blob/master/examples/01_Introducing_txtai.ipynb) |Overview of the functionality provided by txtai||\\n|[Similarity search with images](https://github.com/neuml/txtai/blob/master/examples/13_Similarity_search_with_images.ipynb) |Embed images and text into the same space for search||\\n|[Build a QA database](https://github.com/neuml/txtai/blob/master/examples/34_Build_a_QA_database.ipynb) |Question matching with semantic search||\\n|[Semantic Graphs](https://github.com/neuml/txtai/blob/master/examples/38_Introducing_the_Semantic_Graph.ipynb) |Explore topics, data connectivity and run network analysis||\\n\\n### LLM Orchestration\\nAutonomous agents, retrieval augmented generation (RAG), chat with your data, pipelines and workflows that interface with large language models (LLMs).\\n\\nSee below to learn more.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Prompt templates and task chains](https://github.com/neuml/txtai/blob/master/examples/44_Prompt_templates_and_task_chains.ipynb) |Build model prompts and connect tasks together with workflows||\\n|[Integrate LLM frameworks](https://github.com/neuml/txtai/blob/master/examples/53_Integrate_LLM_Frameworks.ipynb) |Integrate llama.cpp, LiteLLM and custom generation frameworks||\\n|[Build knowledge graphs with LLMs](https://github.'}\n",
|
||||
"{'indexid': 3, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/txtai/blob/master/examples/57_Build_knowledge_graphs_with_LLM_driven_entity_extraction.ipynb) |Build knowledge graphs with LLM-driven entity extraction||\\n\\n#### Agents\\nAgents connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems.\\n\\ntxtai agents are built on top of the Transformers Agent framework. This supports all LLMs txtai supports (Hugging Face, llama.cpp, OpenAI / Claude / AWS Bedrock via LiteLLM).\\n\\nSee the link below to learn more.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Analyzing Hugging Face Posts with Graphs and Agents](https://github.com/neuml/txtai/blob/master/examples/68_Analyzing_Hugging_Face_Posts_with_Graphs_and_Agents.ipynb) |Explore a rich dataset with Graph Analysis and Agents||\\n|[Granting autonomy to agents](https://github.com/neuml/txtai/blob/master/examples/69_Granting_autonomy_to_agents.ipynb) |Agents that iteratively solve problems as they see fit||\\n|[Analyzing LinkedIn Company Posts with Graphs and Agents](https://github.com/neuml/txtai/blob/master/examples/71_Analyzing_LinkedIn_Company_Posts_with_Graphs_and_Agents.ipynb) |Exploring how to improve social media engagement with AI||\\n\\n#### Retrieval augmented generation\\nRetrieval augmented generation (RAG) reduces the risk of LLM hallucinations by constraining the output with a knowledge base as context. RAG is commonly used to \"chat with your data\".\\n\\nA novel feature of txtai is that it can provide both an answer and source citation.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Build RAG pipelines with txtai](https://github.'}\n",
|
||||
"{'indexid': 4, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/txtai/blob/master/examples/52_Build_RAG_pipelines_with_txtai.ipynb) |Guide on retrieval augmented generation including how to create citations||\\n|[How RAG with txtai works](https://github.com/neuml/txtai/blob/master/examples/63_How_RAG_with_txtai_works.ipynb) |Create RAG processes, API services and Docker instances||\\n|[Advanced RAG with graph path traversal](https://github.com/neuml/txtai/blob/master/examples/58_Advanced_RAG_with_graph_path_traversal.ipynb) |Graph path traversal to collect complex sets of data for advanced RAG||\\n|[Speech to Speech RAG](https://github.com/neuml/txtai/blob/master/examples/65_Speech_to_Speech_RAG.ipynb) |Full cycle speech to speech workflow with RAG||\\n\\n### Language Model Workflows\\nLanguage model workflows, also known as semantic workflows, connect language models together to build intelligent applications.\\n\\nWhile LLMs are powerful, there are plenty of smaller, more specialized models that work better and faster for specific tasks. This includes models for extractive question-answering, automatic summarization, text-to-speech, transcription and translation.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Run pipeline workflows](https://github.com/neuml/txtai/blob/master/examples/14_Run_pipeline_workflows.ipynb) |Simple yet powerful constructs to efficiently process data||\\n|[Building abstractive text summaries](https://github.com/neuml/txtai/blob/master/examples/09_Building_abstractive_text_summaries.ipynb) |Run abstractive text summarization||\\n|[Transcribe audio to text](https://github.com/neuml/txtai/blob/master/examples/11_Transcribe_audio_to_text.'}\n",
|
||||
"{'indexid': 5, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'ipynb) |Convert audio files to text||\\n|[Translate text between languages](https://github.com/neuml/txtai/blob/master/examples/12_Translate_text_between_languages.ipynb) |Streamline machine translation and language detection||\\n\\n## Installation\\nThe easiest way to install is via pip and PyPI\\n\\n```\\npip install txtai\\n```\\n\\nPython 3.10+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.\\n\\nSee the detailed [install instructions](https://neuml.github.io/txtai/install) for more information covering [optional dependencies](https://neuml.github.io/txtai/install/#optional-dependencies) , [environment specific prerequisites](https://neuml.github.io/txtai/install/#environment-specific-prerequisites) , [installing from source](https://neuml.github.io/txtai/install/#install-from-source) , [conda support](https://neuml.github.io/txtai/install/#conda) and how to [run with containers](https://neuml.github.io/txtai/cloud) .\\n\\n\\n## Model guide\\nSee the table below for the current recommended models. These models all allow commercial use and offer a blend of speed and performance.\\n\\n|Component|Model(s)|\\n|---|---|\\n|[Embeddings](https://neuml.github.io/txtai/embeddings) |[all-MiniLM-L6-v2](https://hf.co/sentence-transformers/all-MiniLM-L6-v2) |\\n|[Image Captions](https://neuml.github.io/txtai/pipeline/image/caption) |[BLIP](https://hf.co/Salesforce/blip-image-captioning-base) |'}\n",
|
||||
"{'indexid': 5, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'ipynb) |Convert audio files to text||\\n|[Translate text between languages](https://github.com/neuml/txtai/blob/master/examples/12_Translate_text_between_languages.ipynb) |Streamline machine translation and language detection||\\n\\n## Installation\\nThe easiest way to install is via pip and PyPI\\n\\n```\\npip install txtai\\n```\\n\\nPython 3.11+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.\\n\\nSee the detailed [install instructions](https://neuml.github.io/txtai/install) for more information covering [optional dependencies](https://neuml.github.io/txtai/install/#optional-dependencies) , [environment specific prerequisites](https://neuml.github.io/txtai/install/#environment-specific-prerequisites) , [installing from source](https://neuml.github.io/txtai/install/#install-from-source) , [conda support](https://neuml.github.io/txtai/install/#conda) and how to [run with containers](https://neuml.github.io/txtai/cloud) .\\n\\n\\n## Model guide\\nSee the table below for the current recommended models. These models all allow commercial use and offer a blend of speed and performance.\\n\\n|Component|Model(s)|\\n|---|---|\\n|[Embeddings](https://neuml.github.io/txtai/embeddings) |[all-MiniLM-L6-v2](https://hf.co/sentence-transformers/all-MiniLM-L6-v2) |\\n|[Image Captions](https://neuml.github.io/txtai/pipeline/image/caption) |[BLIP](https://hf.co/Salesforce/blip-image-captioning-base) |'}\n",
|
||||
"{'indexid': 6, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '|[Labels - Zero Shot](https://neuml.github.io/txtai/pipeline/text/labels) |[BART-Large-MNLI](https://hf.co/facebook/bart-large) |\\n|[Labels - Fixed](https://neuml.github.io/txtai/pipeline/text/labels) |Fine-tune with [training pipeline](https://neuml.github.io/txtai/pipeline/train/trainer) |\\n|[Large Language Model (LLM)](https://neuml.github.io/txtai/pipeline/text/llm) |[Llama 3.1 Instruct](https://hf.co/meta-llama/Llama-3.1-8B-Instruct) |\\n|[Summarization](https://neuml.github.io/txtai/pipeline/text/summary) |[DistilBART](https://hf.co/sshleifer/distilbart-cnn-12-6) |\\n|[Text-to-Speech](https://neuml.github.io/txtai/pipeline/audio/texttospeech) |[ESPnet JETS](https://hf.co/NeuML/ljspeech-jets-onnx) |\\n|[Transcription](https://neuml.github.io/txtai/pipeline/audio/transcription) |[Whisper](https://hf.co/openai/whisper-base) |\\n|[Translation](https://neuml.github.io/txtai/pipeline/text/translation) |[OPUS Model Series](https://hf.co/Helsinki-NLP) |\\nModels can be loaded as either a path from the Hugging Face Hub or a local directory. Model paths are optional, defaults are loaded when not specified. For tasks with no recommended model, txtai uses the default models as shown in the Hugging Face Tasks guide.\\n\\nSee the following links to learn more.\\n\\n\\n## Powered by txtai\\nThe following applications are powered by txtai.'}\n",
|
||||
"{'indexid': 7, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': \"|Application|Description|\\n|---|---|\\n|[rag](https://github.com/neuml/rag) |Retrieval Augmented Generation (RAG) application|\\n|[ragdata](https://github.com/neuml/ragdata) |Build knowledge bases for RAG|\\n|[paperai](https://github.com/neuml/paperai) |Semantic search and workflows for medical/scientific papers|\\n|[annotateai](https://github.com/neuml/annotateai) |Automatically annotate papers with LLMs|\\nIn addition to this list, there are also many other [open-source projects](https://github.com/neuml/txtai/network/dependents) , [published research](https://scholar.google.com/scholar?q=txtai&hl=en&as_ylo=2022) and closed proprietary/commercial projects that have built on txtai in production.\\n\\n\\n## Further Reading\\n- [Tutorial series on Hashnode](https://neuml.hashnode.dev/series/txtai-tutorial) | [dev.to](https://dev.to/neuml/tutorial-series-on-txtai-ibg) \\n- [What's new in txtai 8.0](https://medium.com/neuml/whats-new-in-txtai-8-0-2d7d0ab4506b) | [7.0](https://medium.com/neuml/whats-new-in-txtai-7-0-855ad6a55440) | [6.0](https://medium.com/neuml/whats-new-in-txtai-6-0-7d93eeedf804) | [5.0](https://medium.com/neuml/whats-new-in-txtai-5-0-e5c75a13b101) | [4.0](https://medium.\"}\n",
|
||||
"{'indexid': 8, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/whats-new-in-txtai-4-0-bbc3a65c3d1c) \\n- [Getting started with semantic search](https://medium.com/neuml/getting-started-with-semantic-search-a9fd9d8a48cf) | [workflows](https://medium.com/neuml/getting-started-with-semantic-workflows-2fefda6165d9) | [rag](https://medium.com/neuml/getting-started-with-rag-9a0cca75f748) \\n\\n## Documentation\\n[Full documentation on txtai](https://neuml.github.io/txtai) including configuration settings for embeddings, pipelines, workflows, API and a FAQ with common questions/issues is available.\\n\\n\\n## Contributing\\nFor those who would like to contribute to txtai, please see [this guide](https://github.com/neuml/.github/blob/master/CONTRIBUTING.md) .'}\n",
|
||||
|
||||
@@ -714,7 +714,7 @@
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Observations: ## Search Results\n",
|
||||
"\n",
|
||||
"|txtai · PyPI](https://pypi.org/project/txtai/)\n",
|
||||
"☁️ Run local or scale out with container orchestration txtai is built with Python 3.10+, Hugging Face Transformers,\n",
|
||||
"☁️ Run local or scale out with container orchestration txtai is built with Python 3.11+, Hugging Face Transformers,\n",
|
||||
"Sentence Transformers and FastAPI. txtai is open-source under an Apache 2.0 license. |!NOTE] NeuML is the company \n",
|
||||
"behind txtai and we provide AI consulting services around our stack. Schedule a meeting or send a message to ...\n",
|
||||
"\n",
|
||||
@@ -767,7 +767,7 @@
|
||||
"Observations: ## Search Results\n",
|
||||
"\n",
|
||||
"|txtai · PyPI](https://pypi.org/project/txtai/)\n",
|
||||
"☁️ Run local or scale out with container orchestration txtai is built with Python 3.10+, Hugging Face Transformers,\n",
|
||||
"☁️ Run local or scale out with container orchestration txtai is built with Python 3.11+, Hugging Face Transformers,\n",
|
||||
"Sentence Transformers and FastAPI. txtai is open-source under an Apache 2.0 license. |!NOTE] NeuML is the company \n",
|
||||
"behind txtai and we provide AI consulting services around our stack. Schedule a meeting or send a message to ...\n",
|
||||
"\n",
|
||||
@@ -896,7 +896,7 @@
|
||||
"|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
|
||||
"- 🔋 Batteries included with defaults to get up and running fast\n",
|
||||
"- ☁️ Run local or scale out with container orchestration\n",
|
||||
"txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
|
||||
"|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
|
||||
"\n",
|
||||
@@ -1060,7 +1060,7 @@
|
||||
"pip install txtai\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"recommended.\n",
|
||||
"\n",
|
||||
"See the detailed |install instructions](https://neuml.github.io/txtai/install) for more information covering \n",
|
||||
@@ -1169,7 +1169,7 @@
|
||||
"|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
|
||||
"- 🔋 Batteries included with defaults to get up and running fast\n",
|
||||
"- ☁️ Run local or scale out with container orchestration\n",
|
||||
"txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
|
||||
"|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
|
||||
"\n",
|
||||
@@ -1333,7 +1333,7 @@
|
||||
"pip install txtai\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"recommended.\n",
|
||||
"\n",
|
||||
"See the detailed |install instructions](https://neuml.github.io/txtai/install) for more information covering \n",
|
||||
@@ -1554,7 +1554,7 @@
|
||||
"|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
|
||||
"- 🔋 Batteries included with defaults to get up and running fast\n",
|
||||
"- ☁️ Run local or scale out with container orchestration\n",
|
||||
"txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
|
||||
"|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
|
||||
"\n",
|
||||
@@ -1599,7 +1599,7 @@
|
||||
"|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
|
||||
"- 🔋 Batteries included with defaults to get up and running fast\n",
|
||||
"- ☁️ Run local or scale out with container orchestration\n",
|
||||
"txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
|
||||
"|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
|
||||
"|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
|
||||
"\n",
|
||||
@@ -1673,7 +1673,7 @@
|
||||
"pip install txtai\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"recommended.\n",
|
||||
"\n",
|
||||
"## Optional dependencies\n",
|
||||
@@ -1847,7 +1847,7 @@
|
||||
"pip install txtai\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
|
||||
"recommended.\n",
|
||||
"\n",
|
||||
"## Optional dependencies\n",
|
||||
@@ -2043,7 +2043,7 @@
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">╭─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮\n",
|
||||
"│ Calling tool: 'write' with arguments: {'path': 'txtai_facts.md', 'content': '# txtai - All-in-One AI │\n",
|
||||
"│ Framework\\n\\n## Overview\\n\\ntxtai is an all-in-one open-source AI framework for semantic search, LLM │\n",
|
||||
"│ orchestration, and language model workflows. It\\'s built with Python 3.10+, Hugging Face Transformers, Sentence │\n",
|
||||
"│ orchestration, and language model workflows. It\\'s built with Python 3.11+, Hugging Face Transformers, Sentence │\n",
|
||||
"│ Transformers, and FastAPI, and is licensed under Apache 2.0.\\n\\n## Key Features\\n\\n- **Vector Search**: │\n",
|
||||
"│ Semantic search with SQL, object storage, topic modeling, graph analysis, and multimodal indexing\\n- │\n",
|
||||
"│ **Embeddings**: Create embeddings for text, documents, audio, images, and video\\n- **Language Model │\n",
|
||||
@@ -2063,7 +2063,7 @@
|
||||
"│ Installation\\n\\nThe easiest way to install is via pip:\\n\\n```bash\\npip install txtai\\n```\\n\\nOptional │\n",
|
||||
"│ dependencies can be installed as extras:\\n\\n```bash\\npip install txtai[all] # Install all dependencies\\npip │\n",
|
||||
"│ install txtai[api] # Serve txtai via a web API\\npip install txtai[workflow] # All workflow tasks\\n```\\n\\n## │\n",
|
||||
"│ Technology Stack\\n\\n- Built with Python 3.10+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n- │\n",
|
||||
"│ Technology Stack\\n\\n- Built with Python 3.11+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n- │\n",
|
||||
"│ FastAPI\\n\\n## Supported Models\\n\\n- Embeddings: all-MiniLM-L6-v2\\n- Image Captions: BLIP\\n- Labels (Zero Shot): │\n",
|
||||
"│ BART-Large-MNLI\\n- Large Language Model (LLM): gpt-oss-20b\\n- Summarization: DistilBART\\n- Text-to-Speech: │\n",
|
||||
"│ ESPnet JETS\\n- Transcription: Whisper\\n- Translation: OPUS Model Series\\n\\n## Company Behind txtai\\n\\nNeuML is │\n",
|
||||
@@ -2081,7 +2081,7 @@
|
||||
"╭─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮\n",
|
||||
"│ Calling tool: 'write' with arguments: {'path': 'txtai_facts.md', 'content': '# txtai - All-in-One AI │\n",
|
||||
"│ Framework\\n\\n## Overview\\n\\ntxtai is an all-in-one open-source AI framework for semantic search, LLM │\n",
|
||||
"│ orchestration, and language model workflows. It\\'s built with Python 3.10+, Hugging Face Transformers, Sentence │\n",
|
||||
"│ orchestration, and language model workflows. It\\'s built with Python 3.11+, Hugging Face Transformers, Sentence │\n",
|
||||
"│ Transformers, and FastAPI, and is licensed under Apache 2.0.\\n\\n## Key Features\\n\\n- **Vector Search**: │\n",
|
||||
"│ Semantic search with SQL, object storage, topic modeling, graph analysis, and multimodal indexing\\n- │\n",
|
||||
"│ **Embeddings**: Create embeddings for text, documents, audio, images, and video\\n- **Language Model │\n",
|
||||
@@ -2101,7 +2101,7 @@
|
||||
"│ Installation\\n\\nThe easiest way to install is via pip:\\n\\n```bash\\npip install txtai\\n```\\n\\nOptional │\n",
|
||||
"│ dependencies can be installed as extras:\\n\\n```bash\\npip install txtai[all] # Install all dependencies\\npip │\n",
|
||||
"│ install txtai[api] # Serve txtai via a web API\\npip install txtai[workflow] # All workflow tasks\\n```\\n\\n## │\n",
|
||||
"│ Technology Stack\\n\\n- Built with Python 3.10+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n- │\n",
|
||||
"│ Technology Stack\\n\\n- Built with Python 3.11+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n- │\n",
|
||||
"│ FastAPI\\n\\n## Supported Models\\n\\n- Embeddings: all-MiniLM-L6-v2\\n- Image Captions: BLIP\\n- Labels (Zero Shot): │\n",
|
||||
"│ BART-Large-MNLI\\n- Large Language Model (LLM): gpt-oss-20b\\n- Summarization: DistilBART\\n- Text-to-Speech: │\n",
|
||||
"│ ESPnet JETS\\n- Transcription: Whisper\\n- Translation: OPUS Model Series\\n\\n## Company Behind txtai\\n\\nNeuML is │\n",
|
||||
@@ -2252,7 +2252,7 @@
|
||||
"\n",
|
||||
"## Overview\n",
|
||||
"\n",
|
||||
"txtai is an all-in-one open-source AI framework for semantic search, LLM orchestration, and language model workflows. It's built with Python 3.10+, Hugging Face Transformers, Sentence Transformers, and FastAPI, and is licensed under Apache 2.0.\n",
|
||||
"txtai is an all-in-one open-source AI framework for semantic search, LLM orchestration, and language model workflows. It's built with Python 3.11+, Hugging Face Transformers, Sentence Transformers, and FastAPI, and is licensed under Apache 2.0.\n",
|
||||
"\n",
|
||||
"## Key Features\n",
|
||||
"\n",
|
||||
|
||||
@@ -109,6 +109,7 @@ nav:
|
||||
- Translation: pipeline/text/translation.md
|
||||
- Train:
|
||||
- HF ONNX: pipeline/train/hfonnx.md
|
||||
- LEMUR Trainer: pipeline/train/lemur.md
|
||||
- ML ONNX: pipeline/train/mlonnx.md
|
||||
- Trainer: pipeline/train/trainer.md
|
||||
- Workflow:
|
||||
|
||||
@@ -20,6 +20,7 @@ default = [
|
||||
"pyyaml>=5.3",
|
||||
"safetensors>=0.4.5",
|
||||
"torch>=2.4",
|
||||
"tqdm>=4.66.3",
|
||||
"transformers>=5.9.0",
|
||||
]
|
||||
|
||||
@@ -52,19 +53,22 @@ extras["dev"] = [
|
||||
"pylint",
|
||||
]
|
||||
|
||||
extras["agent"] = ["jinja2>=3.1.6", "mcpadapt>=0.1.0", "smolagents>=1.23"]
|
||||
extras["agent"] = ["beautifulsoup4>=4.9.3", "jinja2>=3.1.6", "mcp<2.0", "mcpadapt>=0.1.0", "smolagents>=1.23"]
|
||||
|
||||
extras["ann"] = [
|
||||
"annoy>=1.16.3",
|
||||
"bitsandbytes>=0.42.0",
|
||||
"ggml-python>=0.0.41",
|
||||
"hnswlib>=0.5.0",
|
||||
"milvus-lite>=3.1.1",
|
||||
"pgvector>=0.4.1",
|
||||
"rabitqlib>=0.5.1",
|
||||
"scikit-learn>=0.23.1",
|
||||
"scipy>=1.4.1",
|
||||
"sqlalchemy>=2.0.20",
|
||||
"sqlite-vec>=0.1.1",
|
||||
"turbovec>=0.7.0",
|
||||
"zvec>=0.5.1",
|
||||
]
|
||||
|
||||
extras["api"] = [
|
||||
@@ -139,6 +143,7 @@ extras["vectors"] = [
|
||||
"litellm>=1.37.16",
|
||||
"llama-cpp-python>=0.2.75",
|
||||
"model2vec>=0.3.0",
|
||||
"onnxruntime>=1.11.0",
|
||||
"scikit-learn>=0.23.1",
|
||||
"scipy>=1.4.1",
|
||||
"sentence-transformers>=5.0.0",
|
||||
@@ -176,7 +181,7 @@ extras["all"] = (
|
||||
|
||||
setup(
|
||||
name=package,
|
||||
version="9.11.0",
|
||||
version="9.15.0",
|
||||
author="NeuML",
|
||||
description="All-in-one open-source AI framework for semantic search, LLM orchestration and language model workflows",
|
||||
long_description=DESCRIPTION,
|
||||
@@ -191,7 +196,7 @@ setup(
|
||||
packages=find_packages(where="src/python"),
|
||||
package_dir={"": "src/python"},
|
||||
keywords="search embedding machine-learning nlp",
|
||||
python_requires=">=3.10",
|
||||
python_requires=">=3.11",
|
||||
install_requires=install,
|
||||
extras_require=extras,
|
||||
classifiers=[
|
||||
|
||||
@@ -14,12 +14,13 @@ class BashTool(Tool):
|
||||
"""
|
||||
|
||||
# pylint: disable=W0231
|
||||
def __init__(self, allowed=None):
|
||||
def __init__(self, allowed=None, timeout=30):
|
||||
"""
|
||||
Creates a BashTool.
|
||||
|
||||
Args:
|
||||
allowed: list of allowed commands to run, has limited set of defaults which are a best effort not a sandbox
|
||||
timeout: maximum number of seconds to wait for a command to complete, defaults to 30, None disables the timeout
|
||||
"""
|
||||
|
||||
# Tool parameters
|
||||
@@ -31,7 +32,10 @@ class BashTool(Tool):
|
||||
self.output_type = "any"
|
||||
|
||||
# Default list of allowed commands
|
||||
self.allowed = allowed if allowed else ["cat", "cut", "diff", "grep", "head", "ls", "tail"]
|
||||
self.allowed = allowed if allowed is not None else ["cat", "cut", "diff", "grep", "head", "ls", "tail"]
|
||||
|
||||
# Maximum time to wait for a command, None waits indefinitely
|
||||
self.timeout = timeout
|
||||
|
||||
# Validate parameters and initialize tool
|
||||
super().__init__()
|
||||
@@ -41,6 +45,9 @@ class BashTool(Tool):
|
||||
"""
|
||||
Runs a shell command as a subprocess.
|
||||
|
||||
stdin is closed and a timeout is applied. Otherwise an allowed command called with no file arguments
|
||||
(i.e. `cat`) reads from stdin and blocks the agent indefinitely.
|
||||
|
||||
Args:
|
||||
command: command arguments as a list
|
||||
|
||||
@@ -50,6 +57,11 @@ class BashTool(Tool):
|
||||
|
||||
output = None
|
||||
if command and command[0] in self.allowed:
|
||||
output = subprocess.run(command, capture_output=True, text=True, check=False).stdout
|
||||
try:
|
||||
output = subprocess.run(
|
||||
command, capture_output=True, text=True, encoding="utf-8", check=False, stdin=subprocess.DEVNULL, timeout=self.timeout
|
||||
).stdout
|
||||
except subprocess.TimeoutExpired:
|
||||
output = f"Command timed out after {self.timeout} seconds"
|
||||
|
||||
return output
|
||||
|
||||
@@ -32,16 +32,16 @@ class ToolFactory:
|
||||
|
||||
# Default toolkit
|
||||
DEFAULTS = {
|
||||
"bash": BashTool(),
|
||||
"edit": EditTool(),
|
||||
"glob": GlobTool(),
|
||||
"grep": GrepTool(),
|
||||
"python": PythonInterpreterTool(),
|
||||
"question": UserInputTool(),
|
||||
"read": ReadTool(),
|
||||
"todowrite": TodoWriteTool(),
|
||||
"websearch": WebSearchTool(),
|
||||
"write": WriteTool(),
|
||||
"bash": BashTool,
|
||||
"edit": EditTool,
|
||||
"glob": GlobTool,
|
||||
"grep": GrepTool,
|
||||
"python": PythonInterpreterTool,
|
||||
"question": UserInputTool,
|
||||
"read": ReadTool,
|
||||
"todowrite": TodoWriteTool,
|
||||
"websearch": WebSearchTool,
|
||||
"write": WriteTool,
|
||||
}
|
||||
|
||||
# Backwards compatible mappings
|
||||
@@ -81,11 +81,11 @@ class ToolFactory:
|
||||
|
||||
# Get default tool, if applicable
|
||||
elif isinstance(tool, str) and tool in ToolFactory.DEFAULTS:
|
||||
tool = ToolFactory.DEFAULTS[tool]
|
||||
tool = ToolFactory.default(tool)
|
||||
|
||||
# Get ALL default tools, if applicable
|
||||
elif isinstance(tool, str) and tool == "defaults":
|
||||
tools.extend(set(ToolFactory.DEFAULTS.values()))
|
||||
tools.extend(ToolFactory.default(name) for name in ToolFactory.DEFAULTS if name != "webview")
|
||||
tool = None
|
||||
|
||||
# Support importing MCP tool collections
|
||||
@@ -122,6 +122,20 @@ class ToolFactory:
|
||||
except (TypeHintParsingException, TypeError):
|
||||
return ToolFactory.fromdocs(target, config if config else {})
|
||||
|
||||
@staticmethod
|
||||
def default(name):
|
||||
"""
|
||||
Creates a default tool by alias name.
|
||||
|
||||
Args:
|
||||
name: default tool alias name
|
||||
|
||||
Returns:
|
||||
Tool
|
||||
"""
|
||||
|
||||
return ToolFactory.DEFAULTS[name]()
|
||||
|
||||
@staticmethod
|
||||
def fromdocs(target, config):
|
||||
"""
|
||||
|
||||
@@ -6,7 +6,10 @@ from .annoy import Annoy
|
||||
from .factory import ANNFactory
|
||||
from .faiss import Faiss
|
||||
from .hnsw import HNSW
|
||||
from .milvus import Milvus
|
||||
from .numpy import NumPy
|
||||
from .pgvector import PGVector
|
||||
from .rabitq import RabitQ
|
||||
from .torch import Torch
|
||||
from .turbovec import TurboVec
|
||||
from .zvec import Zvec
|
||||
|
||||
@@ -4,15 +4,19 @@ Factory module
|
||||
|
||||
from ...util import Resolver
|
||||
|
||||
from ..base import ANN
|
||||
from .annoy import Annoy
|
||||
from .faiss import Faiss, FAISS
|
||||
from .ggml import GGML
|
||||
from .hnsw import HNSW
|
||||
from .milvus import Milvus
|
||||
from .numpy import NumPy
|
||||
from .pgvector import PGVector
|
||||
from .rabitq import RabitQ
|
||||
from .sqlite import SQLite
|
||||
from .torch import Torch
|
||||
from .turbovec import TurboVec
|
||||
from .zvec import Zvec
|
||||
|
||||
|
||||
class ANNFactory:
|
||||
@@ -32,6 +36,7 @@ class ANNFactory:
|
||||
ANN
|
||||
"""
|
||||
|
||||
# pylint: disable=R0912
|
||||
# ANN instance
|
||||
ann = None
|
||||
backend = config.get("backend", "faiss" if FAISS else "numpy")
|
||||
@@ -41,20 +46,26 @@ class ANNFactory:
|
||||
ann = Annoy(config)
|
||||
elif backend == "faiss":
|
||||
ann = Faiss(config)
|
||||
elif backend == "hnsw":
|
||||
ann = HNSW(config)
|
||||
elif backend == "ggml":
|
||||
ann = GGML(config)
|
||||
elif backend == "hnsw":
|
||||
ann = HNSW(config)
|
||||
elif backend == "milvus":
|
||||
ann = Milvus(config)
|
||||
elif backend == "numpy":
|
||||
ann = NumPy(config)
|
||||
elif backend == "pgvector":
|
||||
ann = PGVector(config)
|
||||
elif backend == "rabitq":
|
||||
ann = RabitQ(config)
|
||||
elif backend == "sqlite":
|
||||
ann = SQLite(config)
|
||||
elif backend == "torch":
|
||||
ann = Torch(config)
|
||||
elif backend == "turbovec":
|
||||
ann = TurboVec(config)
|
||||
elif backend == "zvec":
|
||||
ann = Zvec(config)
|
||||
else:
|
||||
ann = ANNFactory.resolve(backend, config)
|
||||
|
||||
@@ -77,6 +88,6 @@ class ANNFactory:
|
||||
"""
|
||||
|
||||
try:
|
||||
return Resolver()(backend)(config)
|
||||
return Resolver()(backend, ANN)(config)
|
||||
except Exception as e:
|
||||
raise ImportError(f"Unable to resolve ann backend: '{backend}'") from e
|
||||
|
||||
@@ -142,7 +142,7 @@ class Faiss(ANN):
|
||||
|
||||
# Derive quantization. Prefer backend-specific setting. Fallback to root-level parameter.
|
||||
quantize = self.setting("quantize", self.config.get("quantize"))
|
||||
quantize = 8 if isinstance(quantize, bool) else quantize
|
||||
quantize = 8 if isinstance(quantize, bool) and quantize else quantize
|
||||
|
||||
# Get storage setting
|
||||
storage = f"SQ{quantize}" if quantize else "Flat"
|
||||
@@ -229,6 +229,7 @@ class Faiss(ANN):
|
||||
count = self.count()
|
||||
|
||||
default = 6 if count <= 5000 else round(self.cells(count) / 16)
|
||||
default = min(default, getattr(self.backend, "nlist", default))
|
||||
return self.setting("nprobe", default)
|
||||
|
||||
def scores(self, scores):
|
||||
|
||||
@@ -155,7 +155,7 @@ class GGMLTensors:
|
||||
gguf = ggml.gguf_init_from_file(path.encode("utf-8"), params)
|
||||
|
||||
# Load tensors from GGUF file
|
||||
self.loadtensors(context)
|
||||
self.loadtensors(gguf, context)
|
||||
|
||||
# Create graph operation
|
||||
self.graph, self.output = self.creategraph()
|
||||
@@ -218,11 +218,12 @@ class GGMLTensors:
|
||||
ids: ids to delete
|
||||
"""
|
||||
|
||||
shape = utils.get_shape(self.data)
|
||||
rows = self.rows()
|
||||
|
||||
# Filter any index greater than size of array
|
||||
ids = [x for x in ids if x < shape[1]]
|
||||
self.deletes.extend(ids)
|
||||
# Set index ids as deleted, ignoring ids outside the array and ids already marked deleted
|
||||
for x in ids:
|
||||
if x < rows and x not in self.deletes:
|
||||
self.deletes.append(x)
|
||||
|
||||
def search(self, queries):
|
||||
"""
|
||||
@@ -250,8 +251,7 @@ class GGMLTensors:
|
||||
ggml.ggml_backend_graph_compute(self.backend, self.graph)
|
||||
|
||||
# Get size of embeddings data
|
||||
size = utils.get_shape(self.data)
|
||||
size = size[1] if len(size) > 1 else 1
|
||||
size = self.rows()
|
||||
|
||||
# Get and return results
|
||||
results = np.zeros((batch.shape[0], size), dtype=np.float32)
|
||||
@@ -272,7 +272,7 @@ class GGMLTensors:
|
||||
count
|
||||
"""
|
||||
|
||||
return utils.get_shape(self.data)[1] - len(self.deletes) if self.data else 0
|
||||
return self.rows() - len(self.deletes) if self.data else 0
|
||||
|
||||
def save(self, path):
|
||||
"""
|
||||
@@ -288,9 +288,12 @@ class GGMLTensors:
|
||||
# Init and save data tensor
|
||||
gguf = ggml.gguf_init_empty()
|
||||
|
||||
# Add the data tensor
|
||||
ggml.ggml_set_name(self.data, b"data")
|
||||
ggml.gguf_add_tensor(gguf, self.data)
|
||||
# Add the data tensor. GGUF can't store a tensor with zero rows, only save the dimensions for an empty index.
|
||||
if self.rows():
|
||||
ggml.ggml_set_name(self.data, b"data")
|
||||
ggml.gguf_add_tensor(gguf, self.data)
|
||||
else:
|
||||
ggml.gguf_set_val_u32(gguf, b"dimensions", utils.get_shape(self.data)[0])
|
||||
|
||||
# Optionally create and add the deletes tensor
|
||||
if self.deletes:
|
||||
@@ -413,11 +416,12 @@ class GGMLTensors:
|
||||
# Copy embeddings data
|
||||
self.copy(data, self.data, tensortype, 0)
|
||||
|
||||
def loadtensors(self, context):
|
||||
def loadtensors(self, gguf, context):
|
||||
"""
|
||||
Loads existing tensors from context.
|
||||
|
||||
Args:
|
||||
gguf: gguf context
|
||||
context: ggml context
|
||||
"""
|
||||
|
||||
@@ -436,6 +440,10 @@ class GGMLTensors:
|
||||
|
||||
# Copy tensor data to backend
|
||||
ggml.ggml_backend_tensor_set(self.data, ggml.ggml_get_data(data), 0, ggml.ggml_nbytes(data))
|
||||
else:
|
||||
# Empty index, create tensors using the saved dimensions
|
||||
dimensions = ggml.gguf_get_val_u32(gguf, ggml.gguf_find_key(gguf, b"dimensions"))
|
||||
self.createtensors(np.zeros((0, dimensions), dtype=np.float32))
|
||||
|
||||
# Load deletes tensor
|
||||
data = ggml.ggml_get_tensor(context, b"deletes")
|
||||
@@ -465,8 +473,7 @@ class GGMLTensors:
|
||||
queries = ggml.ggml_new_tensor_2d(context, ggml.GGML_TYPE_F32, data.shape[1], self.querysize)
|
||||
|
||||
# Embeddings data with space for both existing and new data
|
||||
shape = utils.get_shape(self.data)
|
||||
merge = ggml.ggml_new_tensor_2d(context, tensortype, data.shape[1], data.shape[0] + shape[1])
|
||||
merge = ggml.ggml_new_tensor_2d(context, tensortype, data.shape[1], data.shape[0] + self.rows())
|
||||
|
||||
# Create new buffer
|
||||
buffer = ggml.ggml_backend_alloc_ctx_tensors(context, backend)
|
||||
@@ -512,7 +519,11 @@ class GGMLTensors:
|
||||
|
||||
# Read tensor type
|
||||
tensortype = self.quantize
|
||||
tensortype = "Q8_0" if isinstance(tensortype, bool) else f"Q{int(tensortype)}_0" if isinstance(tensortype, int) else tensortype
|
||||
tensortype = (
|
||||
"Q8_0"
|
||||
if isinstance(tensortype, bool) and tensortype
|
||||
else f"Q{int(tensortype)}_0" if isinstance(tensortype, int) and not isinstance(tensortype, bool) else tensortype
|
||||
)
|
||||
tensortype = tensortype.upper() if tensortype else "F32"
|
||||
|
||||
# Validate tensor type
|
||||
@@ -572,3 +583,13 @@ class GGMLTensors:
|
||||
"""
|
||||
|
||||
return [queries[x : x + self.querysize] for x in range(0, len(queries), self.querysize)]
|
||||
|
||||
def rows(self):
|
||||
"""
|
||||
Number of rows in the data tensor. Tensors with a single row or no rows report a 1-D shape, so the shape can't be used.
|
||||
|
||||
Returns:
|
||||
row count
|
||||
"""
|
||||
|
||||
return ggml.ggml_nrows(self.data)
|
||||
|
||||
@@ -35,6 +35,9 @@ class HNSW(ANN):
|
||||
self.backend = Index(dim=self.config["dimensions"], space=self.config["metric"])
|
||||
self.backend.load_index(path)
|
||||
|
||||
# Default the delete counter for configurations that don't have it
|
||||
self.config["deletes"] = self.config.get("deletes", 0)
|
||||
|
||||
def index(self, embeddings):
|
||||
# Inner product is equal to cosine similarity on normalized vectors
|
||||
self.config["metric"] = "ip"
|
||||
@@ -85,8 +88,8 @@ class HNSW(ANN):
|
||||
if ef:
|
||||
self.backend.set_ef(ef)
|
||||
|
||||
# Run the query
|
||||
ids, distances = self.backend.knn_query(queries, k=limit)
|
||||
# Run the query, clamp k to the current element count as hnswlib errors when k exceeds it
|
||||
ids, distances = self.backend.knn_query(queries, k=min(limit, self.count()))
|
||||
|
||||
# Map results to [(id, score)]
|
||||
results = []
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
"""
|
||||
milvus-lite module
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
# Conditional import
|
||||
try:
|
||||
import milvus_lite
|
||||
|
||||
MILVUS = True
|
||||
except ImportError:
|
||||
MILVUS = False
|
||||
|
||||
from ...archive import ArchiveFactory
|
||||
from ..base import ANN
|
||||
|
||||
|
||||
class Milvus(ANN):
|
||||
"""
|
||||
Builds an ANN index using the milvus-lite library.
|
||||
"""
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
if not MILVUS:
|
||||
raise ImportError('milvus-lite is not available - install "ann" extra to enable')
|
||||
|
||||
self.collection = None
|
||||
self.directory = None
|
||||
self.path = None
|
||||
|
||||
def load(self, path):
|
||||
self.close()
|
||||
self.directory = tempfile.mkdtemp()
|
||||
self.path = os.path.join(self.directory, "index")
|
||||
|
||||
try:
|
||||
archive = ArchiveFactory.create(self.directory)
|
||||
archive.load(path, "tar")
|
||||
self.open()
|
||||
except Exception:
|
||||
self.close()
|
||||
raise
|
||||
|
||||
def index(self, embeddings):
|
||||
self.close()
|
||||
|
||||
# Lookup index settings
|
||||
m = self.setting("m", 50)
|
||||
|
||||
# Create collection
|
||||
self.directory = tempfile.mkdtemp()
|
||||
self.path = os.path.join(self.directory, "index")
|
||||
|
||||
# Use milvus-lite's native embedded API directly and keep this integration thin
|
||||
self.backend = milvus_lite.MilvusLite(self.path)
|
||||
schema = milvus_lite.CollectionSchema(
|
||||
fields=[
|
||||
milvus_lite.FieldSchema("id", milvus_lite.DataType.INT64, is_primary=True),
|
||||
milvus_lite.FieldSchema("embedding", milvus_lite.DataType.FLOAT_VECTOR, dim=self.config["dimensions"]),
|
||||
]
|
||||
)
|
||||
self.collection = self.backend.create_collection("txtai", schema)
|
||||
self.collection.create_index("embedding", {"index_type": "HNSW", "metric_type": "IP", "params": {"M": m}})
|
||||
self.collection.load()
|
||||
|
||||
# Add items - position in embeddings is used as the id
|
||||
self.insert(embeddings, 0)
|
||||
|
||||
# Add id offset and index build metadata
|
||||
self.config["offset"] = embeddings.shape[0]
|
||||
self.metadata({"m": m, "milvus": milvus_lite.__version__})
|
||||
|
||||
def append(self, embeddings):
|
||||
self.insert(embeddings, self.config["offset"])
|
||||
|
||||
# Update id offset and index metadata
|
||||
self.config["offset"] += embeddings.shape[0]
|
||||
self.metadata()
|
||||
|
||||
def delete(self, ids):
|
||||
# Check length, ids can be a NumPy array
|
||||
if len(ids):
|
||||
self.collection.delete([int(x) for x in ids])
|
||||
|
||||
def search(self, queries, limit):
|
||||
matches = self.collection.search(queries.tolist(), top_k=limit, metric_type="IP", anns_field="embedding")
|
||||
return [[(int(match["id"]), float(match["distance"])) for match in results] for results in matches]
|
||||
|
||||
def count(self):
|
||||
return self.collection.num_entities
|
||||
|
||||
def save(self, path):
|
||||
self.collection.flush()
|
||||
self.backend.close()
|
||||
self.collection = None
|
||||
self.backend = None
|
||||
|
||||
archive = ArchiveFactory.create(self.directory)
|
||||
archive.save(path, "tar")
|
||||
self.open()
|
||||
|
||||
def close(self):
|
||||
# Release milvus-lite collection handles and the directory lock
|
||||
if self.backend:
|
||||
self.backend.close()
|
||||
|
||||
self.collection = None
|
||||
super().close()
|
||||
|
||||
if self.directory and os.path.exists(self.directory):
|
||||
shutil.rmtree(self.directory)
|
||||
|
||||
self.directory = None
|
||||
self.path = None
|
||||
|
||||
def insert(self, embeddings, offset):
|
||||
"""
|
||||
Inserts embeddings starting at offset.
|
||||
|
||||
Args:
|
||||
embeddings: embeddings array
|
||||
offset: starting id
|
||||
"""
|
||||
|
||||
if embeddings.shape[0]:
|
||||
for start in range(0, embeddings.shape[0], 1024):
|
||||
batch = embeddings[start : start + 1024]
|
||||
self.collection.insert([{"id": offset + start + uid, "embedding": embedding.tolist()} for uid, embedding in enumerate(batch)])
|
||||
|
||||
def open(self):
|
||||
"""
|
||||
Opens the milvus-lite database and collection
|
||||
"""
|
||||
|
||||
self.backend = milvus_lite.MilvusLite(self.path)
|
||||
self.collection = self.backend.get_collection("txtai")
|
||||
|
||||
# Reopened collections with an index start released, load to enable search
|
||||
self.collection.load()
|
||||
@@ -60,16 +60,20 @@ class NumPy(ANN):
|
||||
self.metadata()
|
||||
|
||||
def delete(self, ids):
|
||||
# Filter any index greater than size of array
|
||||
ids = [x for x in ids if x < self.backend.shape[0]]
|
||||
# Ignore IDs outside the array; negative IDs must not wrap to live rows.
|
||||
ids = [x for x in ids if 0 <= x < self.backend.shape[0]]
|
||||
|
||||
# Clear specified ids
|
||||
self.backend[ids] = self.tensor(self.zeros((len(ids), self.backend.shape[1])))
|
||||
# Clear specified ids, zeros must match the array data type (e.g. uint8 for quantized data)
|
||||
self.backend[ids] = self.tensor(self.zeros((len(ids), self.backend.shape[1]), dtype=self.backend.dtype))
|
||||
|
||||
def search(self, queries, limit):
|
||||
if self.qbits:
|
||||
# Calculate hamming score for integer vectors
|
||||
scores = self.hammingscore(queries)
|
||||
|
||||
# Zero out deleted (all-zero) rows. Unlike dot product, hamming scores them nonzero,
|
||||
# so they'd otherwise pass the caller's score > 0 filter and resurface in results.
|
||||
scores[:, self.all(self.backend == 0, axis=1)] = 0
|
||||
else:
|
||||
# Dot product on normalized vectors is equal to cosine similarity
|
||||
scores = self.dot(self.tensor(queries), self.backend.T)
|
||||
|
||||
@@ -24,7 +24,6 @@ from ...util import Library
|
||||
np = Library().numpy()
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class PGVector(ANN):
|
||||
"""
|
||||
Builds an ANN index backed by a Postgres database.
|
||||
@@ -72,7 +71,8 @@ class PGVector(ANN):
|
||||
self.metadata()
|
||||
|
||||
def delete(self, ids):
|
||||
self.database.execute(delete(self.table).where(self.table.c["indexid"].in_(ids)))
|
||||
# Bind ids as int, NumPy integers don't match the indexid column
|
||||
self.database.execute(delete(self.table).where(self.table.c["indexid"].in_([int(x) for x in ids])))
|
||||
|
||||
def search(self, queries, limit):
|
||||
results = []
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
"""
|
||||
RabitQ module
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
# Conditional import
|
||||
try:
|
||||
from rabitqlib import HnswIndex, IvfIndex
|
||||
from sklearn.cluster import KMeans
|
||||
|
||||
RABITQ = True
|
||||
except ImportError:
|
||||
RABITQ = False
|
||||
|
||||
from ..base import ANN
|
||||
|
||||
# Core library imports
|
||||
from ...util import Library
|
||||
|
||||
library = Library()
|
||||
np = library.numpy()
|
||||
|
||||
# Sentinel for unfilled top-k slots in native search results
|
||||
SENTINEL = 0xFFFFFFFF
|
||||
|
||||
|
||||
class RabitQ(ANN):
|
||||
"""
|
||||
Builds an ANN index using the rabitqlib library (RaBitQ quantization).
|
||||
|
||||
Only the native index is stored, as a single file. No vectors are kept outside of it. ivf indexes support append
|
||||
and delete natively. The native id of a row is its index id. hnsw indexes can't be extended or have rows removed,
|
||||
so append and delete are not supported in that mode.
|
||||
"""
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
if not RABITQ:
|
||||
raise ImportError('rabitqlib is not available - install "ann" extra to enable')
|
||||
|
||||
# Number of rows removed from the native index (ivf mode only)
|
||||
self.config["deletes"] = self.config.get("deletes", 0)
|
||||
|
||||
def load(self, path):
|
||||
self.close()
|
||||
|
||||
# Invalid modes raise ValueError before native code runs
|
||||
self.backend = IvfIndex.load(path) if self.mode() == "ivf" else HnswIndex.load(path)
|
||||
|
||||
def index(self, embeddings):
|
||||
self.close()
|
||||
|
||||
# Mode is validated before any state is stored
|
||||
mode = self.mode()
|
||||
|
||||
vectors = np.ascontiguousarray(embeddings, dtype=np.float32)
|
||||
rows = vectors.shape[0]
|
||||
|
||||
# Quantization bits per dimension, validated by rabitqlib
|
||||
nbits = self.setting("nbits", 1)
|
||||
|
||||
# Default clusters to 4 * sqrt(N), bounded to [1, N] so KMeans always fits
|
||||
clusters = self.setting("clusters", None)
|
||||
clusters = max(1, min(clusters if clusters else round(4 * math.sqrt(rows)), rows))
|
||||
|
||||
cluster = KMeans(n_clusters=clusters, random_state=0, n_init=1).fit(vectors)
|
||||
centroids = np.ascontiguousarray(cluster.cluster_centers_, dtype=np.float32)
|
||||
clusterids = np.ascontiguousarray(cluster.labels_, dtype=np.uint32)
|
||||
|
||||
dim = self.config["dimensions"]
|
||||
if mode == "ivf":
|
||||
backend = IvfIndex(dim=dim, max_elements=rows, num_clusters=clusters, nbits=nbits, metric="ip")
|
||||
settings = {"mode": mode, "nbits": nbits, "clusters": clusters, "nprobe": self.setting("nprobe", self.nprobe(clusters))}
|
||||
else:
|
||||
backend = HnswIndex(
|
||||
dim=dim,
|
||||
max_elements=rows,
|
||||
M=self.setting("m", 16),
|
||||
ef_construction=self.setting("efconstruction", 200),
|
||||
nbits=nbits,
|
||||
metric="ip",
|
||||
random_seed=self.setting("randomseed", 100),
|
||||
)
|
||||
settings = {
|
||||
"mode": mode,
|
||||
"nbits": nbits,
|
||||
"m": self.setting("m", 16),
|
||||
"efconstruction": self.setting("efconstruction", 200),
|
||||
"efsearch": self.setting("efsearch", None),
|
||||
"randomseed": self.setting("randomseed", 100),
|
||||
}
|
||||
|
||||
backend.build(vectors, centroids, clusterids)
|
||||
self.backend = backend
|
||||
|
||||
# Add id offset and index build metadata
|
||||
self.config["offset"] = rows
|
||||
self.config["deletes"] = 0
|
||||
self.metadata(settings)
|
||||
|
||||
def append(self, embeddings):
|
||||
if self.mode() != "ivf":
|
||||
super().append(embeddings)
|
||||
|
||||
# Native ids continue from the last row ever added
|
||||
embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)
|
||||
self.backend.add(embeddings)
|
||||
|
||||
# Update id offset and index metadata
|
||||
self.config["offset"] = self.backend.max_elements
|
||||
self.metadata()
|
||||
|
||||
def delete(self, ids):
|
||||
if self.mode() != "ivf":
|
||||
super().delete(ids)
|
||||
|
||||
# Native remove hides rows from search. Unknown ids are silently ignored.
|
||||
ids = [uid for uid in ids if 0 <= uid < self.backend.max_elements]
|
||||
if ids:
|
||||
self.config["deletes"] += self.backend.remove(np.array(ids, dtype=np.int64))
|
||||
|
||||
def search(self, queries, limit):
|
||||
if not self.count():
|
||||
return [[] for _ in queries]
|
||||
|
||||
queries = np.ascontiguousarray(queries, dtype=np.float32)
|
||||
|
||||
# The native index errors when k exceeds its row count
|
||||
k = min(limit, self.count())
|
||||
|
||||
if self.mode() == "ivf":
|
||||
nprobe = self.setting("nprobe", self.nprobe(self.backend.num_clusters))
|
||||
ids, distances = self.backend.search(queries, k, nprobe)
|
||||
else:
|
||||
ids, distances = self.backend.search(queries, k, self.setting("efsearch", None) or 0)
|
||||
|
||||
# Map results to [(id, score)], dropping unfilled sentinel slots
|
||||
return [
|
||||
[(int(uid), float(1.0 - dist)) for uid, dist in zip(uids, dists) if uid != SENTINEL]
|
||||
for uids, dists in zip(ids.tolist(), distances.tolist())
|
||||
]
|
||||
|
||||
def count(self):
|
||||
return self.backend.max_elements - self.config["deletes"] if self.backend is not None else 0
|
||||
|
||||
def save(self, path):
|
||||
self.backend.save(path)
|
||||
|
||||
def mode(self):
|
||||
"""
|
||||
Returns the configured index mode.
|
||||
|
||||
Returns:
|
||||
index mode, only "ivf" or "hnsw" are accepted
|
||||
"""
|
||||
|
||||
mode = self.setting("mode", "ivf")
|
||||
if mode not in ("ivf", "hnsw"):
|
||||
raise ValueError(f'Invalid RabitQ mode: {mode}. Supported modes are "ivf" and "hnsw".')
|
||||
|
||||
return mode
|
||||
|
||||
def nprobe(self, clusters):
|
||||
"""
|
||||
Returns the default number of clusters to probe at search time.
|
||||
|
||||
Args:
|
||||
clusters: number of IVF clusters
|
||||
|
||||
Returns:
|
||||
default nprobe
|
||||
"""
|
||||
|
||||
return max(1, round(clusters / 16))
|
||||
@@ -32,9 +32,11 @@ class SQLite(ANN):
|
||||
|
||||
# Quantization setting
|
||||
self.quantize = self.setting("quantize")
|
||||
self.quantize = 8 if isinstance(self.quantize, bool) else int(self.quantize) if self.quantize else None
|
||||
self.quantize = 8 if isinstance(self.quantize, bool) and self.quantize else int(self.quantize) if self.quantize else None
|
||||
|
||||
def load(self, path):
|
||||
if self.path != path:
|
||||
self.close()
|
||||
self.path = path
|
||||
|
||||
def index(self, embeddings):
|
||||
@@ -55,7 +57,8 @@ class SQLite(ANN):
|
||||
self.metadata()
|
||||
|
||||
def delete(self, ids):
|
||||
self.database().executemany(self.deletesql(), [(x,) for x in ids])
|
||||
# Bind ids as int, sqlite3 binds NumPy integers as blobs that match the wrong rows
|
||||
self.database().executemany(self.deletesql(), [(int(x),) for x in ids])
|
||||
|
||||
def search(self, queries, limit):
|
||||
results = []
|
||||
@@ -73,6 +76,9 @@ class SQLite(ANN):
|
||||
return self.cursor.fetchone()[0]
|
||||
|
||||
def save(self, path):
|
||||
# Loading defers opening the connection until the first database operation.
|
||||
self.database()
|
||||
|
||||
# Temporary database
|
||||
if not self.path:
|
||||
# Save temporary database
|
||||
@@ -193,10 +199,13 @@ class SQLite(ANN):
|
||||
connection.execute(self.tablesql())
|
||||
|
||||
# The backup call will hang if there are uncommitted changes, need to copy over
|
||||
# with iterdump (which is much slower)
|
||||
for sql in self.connection.iterdump():
|
||||
if self.tosql('insert into "{table}"') in sql.lower():
|
||||
connection.execute(sql)
|
||||
# the rows (which is much slower). Quantized vectors must be inserted with their type.
|
||||
embedding = "vec_bit(?)" if self.quantize == 1 else "vec_int8(?)" if self.quantize == 8 else "?"
|
||||
rows = self.connection.execute(self.tosql("SELECT indexid, embedding FROM {table}"))
|
||||
connection.executemany(self.tosql(f"INSERT INTO {{table}}(indexid, embedding) VALUES (?, {embedding})"), rows)
|
||||
|
||||
# Commit copied rows, otherwise they are rolled back when the connection is closed
|
||||
connection.commit()
|
||||
else:
|
||||
# Database is up to date, can do a more efficient copy with SQLite C API
|
||||
self.connection.backup(connection)
|
||||
@@ -254,7 +263,12 @@ class SQLite(ANN):
|
||||
SELECT
|
||||
"""
|
||||
|
||||
return self.tosql(("SELECT indexid, 1 - distance FROM {table} " f"WHERE embedding MATCH {self.embeddingsql()} AND k = ? ORDER BY distance"))
|
||||
# BIT distance counts differing bits; normalize it before converting to similarity.
|
||||
# SQLite binary dimensions are the number of bits, not the packed byte count.
|
||||
distance = f"distance / {float(self.config['dimensions'])}" if self.quantize == 1 else "distance"
|
||||
return self.tosql(
|
||||
f"SELECT indexid, 1 - ({distance}) FROM {{table}} " f"WHERE embedding MATCH {self.embeddingsql()} AND k = ? ORDER BY distance"
|
||||
)
|
||||
|
||||
def countsql(self):
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
"""
|
||||
zvec module
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
# Conditional import
|
||||
try:
|
||||
import zvec
|
||||
|
||||
ZVEC = True
|
||||
except ImportError:
|
||||
ZVEC = False
|
||||
|
||||
from ...archive import ArchiveFactory
|
||||
from ..base import ANN
|
||||
|
||||
|
||||
class Zvec(ANN):
|
||||
"""
|
||||
Builds an ANN index using the zvec library.
|
||||
"""
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
if not ZVEC:
|
||||
raise ImportError('zvec is not available - install "ann" extra to enable')
|
||||
|
||||
self.directory = None
|
||||
self.path = None
|
||||
|
||||
def load(self, path):
|
||||
self.close()
|
||||
self.directory = tempfile.mkdtemp()
|
||||
self.path = os.path.join(self.directory, "index")
|
||||
|
||||
try:
|
||||
archive = ArchiveFactory.create(self.directory)
|
||||
archive.load(path, "tar")
|
||||
self.backend = zvec.open(self.path)
|
||||
except Exception:
|
||||
self.close()
|
||||
raise
|
||||
|
||||
def index(self, embeddings):
|
||||
self.close()
|
||||
|
||||
# Lookup index settings
|
||||
efconstruction = self.setting("efconstruction", 200)
|
||||
m = self.setting("m", 50)
|
||||
|
||||
# Create collection
|
||||
self.directory = tempfile.mkdtemp()
|
||||
self.path = os.path.join(self.directory, "index")
|
||||
schema = zvec.CollectionSchema(
|
||||
name="txtai",
|
||||
vectors=zvec.VectorSchema(
|
||||
"embedding",
|
||||
self.datatype(),
|
||||
self.config["dimensions"],
|
||||
index_param=zvec.HnswIndexParam(metric_type=zvec.MetricType.IP, m=m, ef_construction=efconstruction),
|
||||
),
|
||||
)
|
||||
self.backend = zvec.create_and_open(path=self.path, schema=schema)
|
||||
|
||||
# Add items - position in embeddings is used as the id
|
||||
self.insert(embeddings, 0)
|
||||
|
||||
# Build HNSW index
|
||||
self.backend.optimize()
|
||||
|
||||
# Add id offset and index build metadata
|
||||
self.config["offset"] = embeddings.shape[0]
|
||||
self.metadata({"efconstruction": efconstruction, "m": m, "zvec": zvec.__version__})
|
||||
|
||||
def append(self, embeddings):
|
||||
self.insert(embeddings, self.config["offset"])
|
||||
|
||||
# Add new items to HNSW index
|
||||
self.backend.optimize()
|
||||
|
||||
# Update id offset and index metadata
|
||||
self.config["offset"] += embeddings.shape[0]
|
||||
self.metadata()
|
||||
|
||||
def delete(self, ids):
|
||||
# Check length, ids can be a NumPy array
|
||||
if len(ids):
|
||||
self.backend.delete([str(uid) for uid in ids])
|
||||
|
||||
def search(self, queries, limit):
|
||||
# Lookup search settings
|
||||
param = zvec.HnswQueryParam(ef=self.setting("efsearch", 300))
|
||||
|
||||
results = []
|
||||
for query in queries:
|
||||
matches = self.backend.query(
|
||||
zvec.Query(field_name="embedding", vector=self.prepare(query), param=param),
|
||||
topk=limit,
|
||||
)
|
||||
results.append([(int(match.id), float(match.score)) for match in matches])
|
||||
|
||||
return results
|
||||
|
||||
def count(self):
|
||||
return self.backend.stats.doc_count
|
||||
|
||||
def save(self, path):
|
||||
self.backend.flush()
|
||||
self.backend = None
|
||||
|
||||
archive = ArchiveFactory.create(self.directory)
|
||||
archive.save(path, "tar")
|
||||
self.backend = zvec.open(self.path)
|
||||
|
||||
def close(self):
|
||||
# Parent logic releases the collection and its file locks
|
||||
super().close()
|
||||
|
||||
if self.directory and os.path.exists(self.directory):
|
||||
shutil.rmtree(self.directory)
|
||||
|
||||
self.directory = None
|
||||
self.path = None
|
||||
|
||||
def insert(self, embeddings, offset):
|
||||
"""
|
||||
Inserts embeddings starting at offset.
|
||||
|
||||
Args:
|
||||
embeddings: embeddings array
|
||||
offset: starting id
|
||||
"""
|
||||
|
||||
if embeddings.shape[0]:
|
||||
for start in range(0, embeddings.shape[0], 1024):
|
||||
batch = embeddings[start : start + 1024]
|
||||
self.backend.insert(
|
||||
[zvec.Doc(id=str(offset + start + uid), vectors={"embedding": self.prepare(embedding)}) for uid, embedding in enumerate(batch)]
|
||||
)
|
||||
|
||||
def datatype(self):
|
||||
"""
|
||||
Gets the vector data type.
|
||||
|
||||
Returns:
|
||||
zvec vector data type
|
||||
"""
|
||||
|
||||
return zvec.DataType.VECTOR_FP32
|
||||
|
||||
def prepare(self, data):
|
||||
"""
|
||||
Prepares data for the vector field.
|
||||
|
||||
Args:
|
||||
data: input data
|
||||
|
||||
Returns:
|
||||
data formatted for zvec
|
||||
"""
|
||||
|
||||
return data.tolist()
|
||||
@@ -5,3 +5,4 @@ Sparse ANN imports
|
||||
from .factory import SparseANNFactory
|
||||
from .ivfsparse import IVFSparse
|
||||
from .pgsparse import PGSparse
|
||||
from .zvecsparse import ZvecSparse
|
||||
|
||||
@@ -4,8 +4,10 @@ Factory module
|
||||
|
||||
from ...util import Resolver
|
||||
|
||||
from ..base import ANN
|
||||
from .ivfsparse import IVFSparse
|
||||
from .pgsparse import PGSparse
|
||||
from .zvecsparse import ZvecSparse
|
||||
|
||||
|
||||
class SparseANNFactory:
|
||||
@@ -34,6 +36,8 @@ class SparseANNFactory:
|
||||
ann = IVFSparse(config)
|
||||
elif backend == "pgsparse":
|
||||
ann = PGSparse(config)
|
||||
elif backend == "zvecsparse":
|
||||
ann = ZvecSparse(config)
|
||||
else:
|
||||
ann = SparseANNFactory.resolve(backend, config)
|
||||
|
||||
@@ -56,6 +60,6 @@ class SparseANNFactory:
|
||||
"""
|
||||
|
||||
try:
|
||||
return Resolver()(backend)(config)
|
||||
return Resolver()(backend, ANN)(config)
|
||||
except Exception as e:
|
||||
raise ImportError(f"Unable to resolve sparse ann backend: '{backend}'") from e
|
||||
|
||||
@@ -115,8 +115,12 @@ class IVFSparse(ANN):
|
||||
self.metadata()
|
||||
|
||||
def delete(self, ids):
|
||||
# Set index ids as deleted
|
||||
self.deletes.extend(ids)
|
||||
# Set index ids as deleted, ignoring ids that were never indexed and ids already marked deleted
|
||||
# Ids are stored as int, the deletes list is serialized with msgpack
|
||||
size = self.size()
|
||||
for x in ids:
|
||||
if x < size and x not in self.deletes:
|
||||
self.deletes.append(int(x))
|
||||
|
||||
def search(self, queries, limit):
|
||||
results = []
|
||||
@@ -204,10 +208,10 @@ class IVFSparse(ANN):
|
||||
cluster centroids
|
||||
"""
|
||||
|
||||
# Select top n most important features that contribute to L2 vector norm
|
||||
indices = np.argsort(-norm(train, axis=0))[: self.setting("nfeatures", 25)]
|
||||
data = train[:, indices]
|
||||
data = train
|
||||
# Select top n most important features that contribute to L2 vector norm.
|
||||
# nfeatures is opt-in: the full feature space is used when it isn't set.
|
||||
nfeatures = self.setting("nfeatures")
|
||||
data = train[:, np.argsort(-norm(train, axis=0))[:nfeatures]] if nfeatures else train
|
||||
|
||||
# Cluster data using k-means
|
||||
kmeans = MiniBatchKMeans(n_clusters=clusters, random_state=0, n_init=5).fit(data)
|
||||
@@ -273,6 +277,10 @@ class IVFSparse(ANN):
|
||||
indices = np.argpartition(-scores, limit if limit < scores.shape[1] else scores.shape[1] - 1)[:, :limit]
|
||||
scores = np.take_along_axis(scores, indices, axis=1)
|
||||
|
||||
# argpartition doesn't order within the top n - sort each row by score descending
|
||||
order = np.argsort(-scores, axis=1)
|
||||
indices, scores = np.take_along_axis(indices, order, axis=1), np.take_along_axis(scores, order, axis=1)
|
||||
|
||||
return indices, scores
|
||||
|
||||
def scan(self, query, limit, blockids):
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
ZvecSparse module
|
||||
"""
|
||||
|
||||
# Conditional import
|
||||
try:
|
||||
import zvec
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
from ..dense import Zvec
|
||||
|
||||
|
||||
class ZvecSparse(Zvec):
|
||||
"""
|
||||
Builds a Sparse ANN index using the zvec library.
|
||||
"""
|
||||
|
||||
def search(self, queries, limit):
|
||||
# Lookup search settings
|
||||
param = zvec.HnswQueryParam(ef=self.setting("efsearch", 300))
|
||||
|
||||
results = []
|
||||
for query in queries:
|
||||
# zvec requires at least one non-zero value per query
|
||||
matches = self.backend.query(zvec.Query(field_name="embedding", vector=self.prepare(query), param=param), topk=limit) if query.nnz else []
|
||||
results.append([(int(match.id), float(match.score)) for match in matches])
|
||||
|
||||
return results
|
||||
|
||||
def datatype(self):
|
||||
return zvec.DataType.SPARSE_VECTOR_FP32
|
||||
|
||||
def prepare(self, data):
|
||||
return {int(index): float(value) for index, value in zip(data.indices, data.data)}
|
||||
@@ -12,6 +12,7 @@ from httpx import AsyncClient, ASGITransport
|
||||
|
||||
from .authorization import Authorization
|
||||
from .base import API
|
||||
from .extension import Extension
|
||||
from .factory import APIFactory
|
||||
|
||||
from ..app import Application
|
||||
@@ -95,15 +96,15 @@ def lifespan(application):
|
||||
|
||||
# Conditionally add routes based on configuration
|
||||
for name, router in routers.items():
|
||||
if name in config:
|
||||
if name in config and enabled(config, name):
|
||||
application.include_router(router)
|
||||
|
||||
# Special case for embeddings clusters
|
||||
if "cluster" in config and "embeddings" not in config:
|
||||
if "cluster" in config and "embeddings" not in config and enabled(config, "embeddings"):
|
||||
application.include_router(routers["embeddings"])
|
||||
|
||||
# Special case to add similarity instance for embeddings
|
||||
if "embeddings" in config and "similarity" not in config:
|
||||
if "embeddings" in config and "similarity" not in config and enabled(config, "similarity"):
|
||||
application.include_router(routers["similarity"])
|
||||
|
||||
# Execute extensions if present
|
||||
@@ -111,7 +112,7 @@ def lifespan(application):
|
||||
if extensions:
|
||||
for extension in extensions.split(","):
|
||||
# Create instance and execute extension
|
||||
extension = APIFactory.get(extension.strip())()
|
||||
extension = APIFactory.get(extension.strip(), Extension)()
|
||||
extension(application)
|
||||
|
||||
# Add Model Context Protocol (MCP) service, if applicable
|
||||
@@ -120,6 +121,21 @@ def lifespan(application):
|
||||
yield
|
||||
|
||||
|
||||
def enabled(config, name):
|
||||
"""
|
||||
Check configuration to determine if a route is enabled or disabled
|
||||
|
||||
Args:
|
||||
config: configuration
|
||||
name: route name
|
||||
|
||||
Returns:
|
||||
True if a route is enabled, False otherwise
|
||||
"""
|
||||
|
||||
return config.get("routes", {}).get(name, True)
|
||||
|
||||
|
||||
def createmcp(application, config):
|
||||
"""
|
||||
Create a MCP service if necessary.
|
||||
|
||||
@@ -3,9 +3,10 @@ Authorization module
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import hmac
|
||||
|
||||
from fastapi import Header, HTTPException
|
||||
from fastapi import Depends, HTTPException
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
|
||||
|
||||
class Authorization:
|
||||
@@ -13,17 +14,17 @@ class Authorization:
|
||||
Basic token authorization.
|
||||
"""
|
||||
|
||||
def __init__(self, token=None):
|
||||
def __init__(self, token):
|
||||
"""
|
||||
Creates a new Authorization instance.
|
||||
|
||||
Args:
|
||||
token: SHA-256 hash of token to check
|
||||
token: SHA-256 hash of the valid authorization token
|
||||
"""
|
||||
|
||||
self.token = token if token else os.environ.get("TOKEN")
|
||||
self.token = token
|
||||
|
||||
def __call__(self, authorization: str = Header(default=None)):
|
||||
def __call__(self, authorization: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
|
||||
"""
|
||||
Validates authorization header is present and equal to current token.
|
||||
|
||||
@@ -31,23 +32,19 @@ class Authorization:
|
||||
authorization: authorization header
|
||||
"""
|
||||
|
||||
if not authorization or self.token != self.digest(authorization):
|
||||
# Authorization header is parsed into {scheme="Bearer", credentials="<token>"}
|
||||
if not hmac.compare_digest(self.token, self.digest(authorization.credentials)):
|
||||
raise HTTPException(status_code=401, detail="Invalid Authorization Token")
|
||||
|
||||
def digest(self, authorization):
|
||||
def digest(self, token):
|
||||
"""
|
||||
Computes a SHA-256 hash for input authorization token.
|
||||
|
||||
Args:
|
||||
authorization: authorization header
|
||||
token: authorization token
|
||||
|
||||
Returns:
|
||||
SHA-256 hash of authorization token
|
||||
"""
|
||||
|
||||
# Replace Bearer prefix
|
||||
prefix = "Bearer "
|
||||
token = authorization[len(prefix) :] if authorization.startswith(prefix) else authorization
|
||||
|
||||
# Compute SHA-256 hash
|
||||
return hashlib.sha256(token.encode("utf-8")).hexdigest()
|
||||
|
||||
@@ -32,7 +32,7 @@ class API(Application):
|
||||
weights = self.weights(request.query_params.get("weights") if request and hasattr(request, "query_params") else weights)
|
||||
index = request.query_params.get("index") if request and hasattr(request, "query_params") else index
|
||||
parameters = request.query_params.get("parameters") if request and hasattr(request, "query_params") else parameters
|
||||
graph = request.query_params.get("graph") if request and hasattr(request, "query_params") else graph
|
||||
graph = self.graph(request.query_params.get("graph") if request and hasattr(request, "query_params") else graph)
|
||||
|
||||
# Decode parameters
|
||||
parameters = json.loads(parameters) if parameters and isinstance(parameters, str) else parameters
|
||||
@@ -157,3 +157,17 @@ class API(Application):
|
||||
"""
|
||||
|
||||
return float(weights) if weights else weights
|
||||
|
||||
def graph(self, graph):
|
||||
"""
|
||||
Parses the graph parameter from the request.
|
||||
|
||||
Args:
|
||||
graph: graph parameter
|
||||
|
||||
Returns:
|
||||
graph as a bool
|
||||
"""
|
||||
|
||||
# Query parameters are strings - accept the same values the batchsearch endpoint coerces
|
||||
return graph.strip().lower() in ("1", "true", "yes", "on") if isinstance(graph, str) else bool(graph)
|
||||
|
||||
@@ -56,17 +56,19 @@ class Cluster:
|
||||
"""
|
||||
|
||||
# Build URL
|
||||
action = f"search?query={urllib.parse.quote_plus(query)}"
|
||||
params = {"query": query}
|
||||
if limit:
|
||||
action += f"&limit={limit}"
|
||||
if weights:
|
||||
action += f"&weights={weights}"
|
||||
params["limit"] = limit
|
||||
if weights is not None:
|
||||
params["weights"] = weights
|
||||
if index:
|
||||
action += f"&index={index}"
|
||||
params["index"] = index
|
||||
if parameters:
|
||||
action += f"¶meters={json.dumps(parameters) if isinstance(parameters, dict) else parameters}"
|
||||
params["parameters"] = json.dumps(parameters) if isinstance(parameters, dict) else parameters
|
||||
if graph is not None:
|
||||
action += f"&graph={graph}"
|
||||
params["graph"] = graph
|
||||
|
||||
action = f"search?{urllib.parse.urlencode(params)}"
|
||||
|
||||
# Run query and flatten results into single results list
|
||||
results = []
|
||||
@@ -100,7 +102,7 @@ class Cluster:
|
||||
params = {"queries": queries}
|
||||
if limit:
|
||||
params["limit"] = limit
|
||||
if weights:
|
||||
if weights is not None:
|
||||
params["weights"] = weights
|
||||
if index:
|
||||
params["index"] = index
|
||||
@@ -196,11 +198,11 @@ class Cluster:
|
||||
shards = [[] for _ in range(len(self.shards))]
|
||||
for document in documents:
|
||||
uid = document.get("id") if isinstance(document, dict) else document
|
||||
if uid and isinstance(uid, str):
|
||||
# Quick int hash of string to help derive shard id
|
||||
if isinstance(uid, str):
|
||||
# Quick int hash of string to help derive shard id (handles empty strings too)
|
||||
uid = zlib.adler32(uid.encode("utf-8"))
|
||||
elif uid is None:
|
||||
# Get random shard id when uid isn't set
|
||||
elif not isinstance(uid, int):
|
||||
# Get random shard id when uid isn't a usable int (None, a tuple document, etc.)
|
||||
uid = random.randint(0, len(shards) - 1)
|
||||
|
||||
shards[uid % len(self.shards)].append(document)
|
||||
|
||||
@@ -4,6 +4,8 @@ API factory module
|
||||
|
||||
from ..util import Resolver
|
||||
|
||||
from .base import API
|
||||
|
||||
|
||||
class APIFactory:
|
||||
"""
|
||||
@@ -11,18 +13,19 @@ class APIFactory:
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def get(api):
|
||||
def get(path, base=None):
|
||||
"""
|
||||
Gets a new instance of api class.
|
||||
Gets a new instance of a class.
|
||||
|
||||
Args:
|
||||
api: API instance class
|
||||
path: class path
|
||||
base: optional required base class
|
||||
|
||||
Returns:
|
||||
API
|
||||
Class
|
||||
"""
|
||||
|
||||
return Resolver()(api)
|
||||
return Resolver()(path, base)
|
||||
|
||||
@staticmethod
|
||||
def create(config, api):
|
||||
@@ -37,4 +40,4 @@ class APIFactory:
|
||||
API instance
|
||||
"""
|
||||
|
||||
return APIFactory.get(api)(config)
|
||||
return APIFactory.get(api, API)(config)
|
||||
|
||||
@@ -4,6 +4,7 @@ Route module
|
||||
|
||||
import json
|
||||
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi.routing import APIRoute
|
||||
|
||||
from .responses import ResponseFactory, MessagePackResponse
|
||||
@@ -40,7 +41,7 @@ class EncodingAPIRoute(APIRoute):
|
||||
status_code=response.status_code,
|
||||
headers={k: v for k, v in response.headers.items() if k.lower() not in ["content-length", "content-type"]},
|
||||
)
|
||||
if not isinstance(response, target) and target == MessagePackResponse
|
||||
if not isinstance(response, (target, StreamingResponse)) and target == MessagePackResponse
|
||||
else response
|
||||
)
|
||||
|
||||
|
||||
@@ -44,8 +44,8 @@ def chat(
|
||||
# Build keyword arguments
|
||||
kwargs = {key: value for key, value in [("stream", stream), ("maxlength", max_completion_tokens)] if value}
|
||||
|
||||
# Get first message
|
||||
message = messages[0]["content"]
|
||||
# Get latest message
|
||||
message = messages[-1]["content"]
|
||||
|
||||
# Agent
|
||||
if model in application.get().agents:
|
||||
@@ -165,9 +165,9 @@ class ChatResponse:
|
||||
return {
|
||||
"id": str(uuid.uuid4()),
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time() * 1000),
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{"id": 0, "message": {"role": "assistant", "content": result}, "finish_reason": "stop"}],
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": result}, "finish_reason": "stop"}],
|
||||
}
|
||||
|
||||
|
||||
@@ -177,15 +177,38 @@ class StreamingChatResponse:
|
||||
"""
|
||||
|
||||
def __call__(self, model, result):
|
||||
# All chunks of a response share the same id and creation time
|
||||
uid, created = str(uuid.uuid4()), int(time.time())
|
||||
|
||||
for chunk in result:
|
||||
yield "data: " + json.dumps(
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time() * 1000),
|
||||
"model": model,
|
||||
"choices": [{"id": 0, "delta": {"content": chunk}}],
|
||||
}
|
||||
) + "\n\n"
|
||||
yield self.chunk(uid, created, model, {"content": chunk}, None)
|
||||
|
||||
# Final chunk signals the end of the response
|
||||
yield self.chunk(uid, created, model, {}, "stop")
|
||||
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
def chunk(self, uid, created, model, delta, reason):
|
||||
"""
|
||||
Builds a server-sent event for a single chat completion chunk.
|
||||
|
||||
Args:
|
||||
uid: response id
|
||||
created: response creation time in seconds
|
||||
model: model name
|
||||
delta: chunk content
|
||||
reason: finish reason, None until the final chunk
|
||||
|
||||
Returns:
|
||||
server-sent event
|
||||
"""
|
||||
|
||||
data = {
|
||||
"id": uid,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": reason}],
|
||||
}
|
||||
|
||||
return "data: " + json.dumps(data) + "\n\n"
|
||||
|
||||
@@ -18,7 +18,6 @@ from ..util import Library
|
||||
yaml = Library().yaml()
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class Application:
|
||||
"""
|
||||
Builds YAML-configured txtai applications.
|
||||
@@ -259,7 +258,7 @@ class Application:
|
||||
"""
|
||||
|
||||
# Check for task shorthand syntax
|
||||
task = {"action": task} if isinstance(task, (str, list)) else task
|
||||
task = {"action": task} if isinstance(task, (str, list)) else task.copy()
|
||||
|
||||
if "action" in task:
|
||||
action = task["action"]
|
||||
@@ -291,6 +290,10 @@ class Application:
|
||||
# Save resolved action(s)
|
||||
task["action"] = actions[0] if not isinstance(action, list) else actions
|
||||
|
||||
if "task" in task and task["task"] == "retrieve":
|
||||
# Default to safeopen enabled
|
||||
task["safeopen"] = task.get("safeopen", True)
|
||||
|
||||
# Resolve initializer
|
||||
if "initialize" in task and isinstance(task["initialize"], str):
|
||||
task["initialize"] = self.function(task["initialize"])
|
||||
@@ -350,6 +353,9 @@ class Application:
|
||||
resolved function
|
||||
"""
|
||||
|
||||
# Function configuration
|
||||
config = {}
|
||||
|
||||
# Check if function is a pipeline
|
||||
if function in self.pipelines:
|
||||
return self.pipelines[function]
|
||||
@@ -358,8 +364,12 @@ class Application:
|
||||
if function in self.workflows:
|
||||
return self.workflows[function]
|
||||
|
||||
if function in ("textractor", "urlretrieve"):
|
||||
# Default to safeopen enabled
|
||||
config["safeopen"] = config.get("safeopen", True)
|
||||
|
||||
# Attempt to resolve action as a callable function
|
||||
return PipelineFactory.create({}, function)
|
||||
return PipelineFactory.create(config, function)
|
||||
|
||||
def search(self, query, limit=10, weights=None, index=None, parameters=None, graph=False):
|
||||
"""
|
||||
|
||||
@@ -4,6 +4,7 @@ Factory module
|
||||
|
||||
from ..util import Resolver
|
||||
|
||||
from .base import Cloud
|
||||
from .hub import HuggingFaceHub
|
||||
from .storage import ObjectStorage, LIBCLOUD
|
||||
|
||||
@@ -58,7 +59,7 @@ class CloudFactory:
|
||||
"""
|
||||
|
||||
try:
|
||||
return Resolver()(backend)(config)
|
||||
return Resolver()(backend, Cloud)(config)
|
||||
|
||||
except Exception as e:
|
||||
# Failure message
|
||||
|
||||
@@ -154,6 +154,11 @@ class Console(Cmd):
|
||||
else:
|
||||
results = self.app.search(query, limit=self.vlimit)
|
||||
|
||||
# No matches, nothing to render
|
||||
if not results:
|
||||
self.console.print("No results")
|
||||
return
|
||||
|
||||
columns, table = {}, Table(box=box.SQUARE, style="#03a9f4")
|
||||
|
||||
# Build column list
|
||||
|
||||
@@ -24,7 +24,7 @@ class Data:
|
||||
self.columns = columns
|
||||
self.maxlength = maxlength
|
||||
|
||||
def __call__(self, train, validation, workers):
|
||||
def __call__(self, train, validation, workers, batch):
|
||||
"""
|
||||
Tokenizes training and validation data and returns processed datasets.
|
||||
|
||||
@@ -32,14 +32,18 @@ class Data:
|
||||
train: training data
|
||||
validation: validation data
|
||||
workers: number of concurrent tokenizers when processing datasets, only main process used when set to None
|
||||
batch: tokenize batch size
|
||||
|
||||
Returns:
|
||||
(train, validation)
|
||||
"""
|
||||
|
||||
return (self.prepare(train, self.process, workers), self.prepare(validation, self.process, workers) if validation else None)
|
||||
return (
|
||||
self.prepare(train, self.process, workers, batch),
|
||||
self.prepare(validation, self.process, workers, batch) if validation is not None else None,
|
||||
)
|
||||
|
||||
def prepare(self, data, fn, workers):
|
||||
def prepare(self, data, fn, workers, batch):
|
||||
"""
|
||||
Prepares and tokenizes data for training.
|
||||
|
||||
@@ -47,14 +51,16 @@ class Data:
|
||||
data: input data
|
||||
fn: tokenize processing function to apply
|
||||
workers: number of concurrent tokenizers when processing datasets, only main process used when set to None
|
||||
batch: tokenize batch size
|
||||
|
||||
Returns:
|
||||
tokens
|
||||
"""
|
||||
|
||||
if hasattr(data, "map"):
|
||||
# Check for column_names, pandas DataFrames also have a map method
|
||||
if hasattr(data, "column_names"):
|
||||
# Hugging Face dataset
|
||||
tokens = data.map(fn, batched=True, num_proc=workers, remove_columns=data.column_names)
|
||||
tokens = data.map(fn, batched=True, batch_size=batch, num_proc=workers, remove_columns=data.column_names)
|
||||
else:
|
||||
# Re-orient data into columns for efficient batch tokenization
|
||||
columns = {}
|
||||
@@ -91,11 +97,11 @@ class Data:
|
||||
column = self.columns[-1]
|
||||
|
||||
# Return length of labels if it's an array
|
||||
length = self.length(data[column][0] if hasattr(data, "columns") else data[0][column])
|
||||
length = self.length(next(iter(data[column])) if hasattr(data, "columns") else data[0][column])
|
||||
if length:
|
||||
return length
|
||||
|
||||
if hasattr(data, "map"):
|
||||
if hasattr(data, "column_names"):
|
||||
# Hugging Face dataset
|
||||
labels = sorted(data.unique(self.columns[-1]))
|
||||
elif hasattr(data, "columns"):
|
||||
|
||||
@@ -72,14 +72,20 @@ class Questions(Data):
|
||||
while sequences[end] != (1 if self.rpad else 0):
|
||||
end -= 1
|
||||
|
||||
# Map start character and end character to matching token index
|
||||
while start < len(offset) and offset[start][0] <= startchar:
|
||||
start += 1
|
||||
tokenized["start_positions"].append(start - 1)
|
||||
# Label with CLS token if the answer isn't fully in this span. Long contexts are split into
|
||||
# overlapping spans with stride and most spans don't contain the answer.
|
||||
if not (offset[start][0] <= startchar and offset[end][1] >= endchar):
|
||||
tokenized["start_positions"].append(clstoken)
|
||||
tokenized["end_positions"].append(clstoken)
|
||||
else:
|
||||
# Map start character and end character to matching token index
|
||||
while start < len(offset) and offset[start][0] <= startchar:
|
||||
start += 1
|
||||
tokenized["start_positions"].append(start - 1)
|
||||
|
||||
while offset[end][1] >= endchar:
|
||||
end -= 1
|
||||
tokenized["end_positions"].append(end + 1)
|
||||
while offset[end][1] >= endchar:
|
||||
end -= 1
|
||||
tokenized["end_positions"].append(end + 1)
|
||||
|
||||
return tokenized
|
||||
|
||||
|
||||
@@ -104,7 +104,7 @@ class Texts(Data):
|
||||
nextlength = len(rows[x + 1]) if x < len(rows) - 1 else 0
|
||||
|
||||
# New row
|
||||
if (length + nextlength) >= self.maxlength:
|
||||
if (length + nextlength) > self.maxlength:
|
||||
for column in columns:
|
||||
results[column].append(list(chain(*inputs[column][index : x + 1])))
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ except ImportError:
|
||||
|
||||
from .embedded import Embedded
|
||||
from .schema import Statement
|
||||
from .sql import SQLError
|
||||
|
||||
|
||||
class DuckDB(Embedded):
|
||||
@@ -29,6 +30,11 @@ class DuckDB(Embedded):
|
||||
DELETE_DOCUMENT = "DELETE FROM documents WHERE id = ?"
|
||||
DELETE_OBJECT = "DELETE FROM objects WHERE id = ?"
|
||||
|
||||
# Added to errors that abort the transaction
|
||||
ABORTED = (
|
||||
"This error aborted the DuckDB transaction. Unsaved changes are lost and later operations fail until the index is loaded or built again."
|
||||
)
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
@@ -36,8 +42,15 @@ class DuckDB(Embedded):
|
||||
raise ImportError('DuckDB is not available - install "database" extra to enable')
|
||||
|
||||
def execute(self, function, *args):
|
||||
# Call parent method with DuckDB compatible arguments
|
||||
return super().execute(function, *self.formatargs(args))
|
||||
try:
|
||||
# Call parent method with DuckDB compatible arguments
|
||||
return super().execute(function, *self.formatargs(args))
|
||||
except SQLError as e:
|
||||
# Runtime errors abort the transaction. DuckDB has no savepoints, so it can't be recovered without a rollback.
|
||||
if self.aborted():
|
||||
raise SQLError(f"{e}\n\n{DuckDB.ABORTED}") from None
|
||||
|
||||
raise
|
||||
|
||||
def insertdocument(self, uid, data, tags, entry):
|
||||
# Delete existing document
|
||||
@@ -123,6 +136,20 @@ class DuckDB(Embedded):
|
||||
|
||||
return connection
|
||||
|
||||
def aborted(self):
|
||||
"""
|
||||
Checks if the current transaction is aborted.
|
||||
|
||||
Returns:
|
||||
True if the current transaction is aborted
|
||||
"""
|
||||
|
||||
try:
|
||||
self.connection.execute("SELECT 1")
|
||||
return False
|
||||
except duckdb.TransactionException:
|
||||
return True
|
||||
|
||||
def formatargs(self, args):
|
||||
"""
|
||||
DuckDB doesn't support named parameters. This method replaces named parameters with question marks
|
||||
@@ -139,19 +166,16 @@ class DuckDB(Embedded):
|
||||
# Unpack query args
|
||||
query, parameters = args
|
||||
|
||||
# Iterate over parameters
|
||||
# - Replace named parameters with ?'s
|
||||
# - Build list of value with position indexes
|
||||
# Match all named parameters in a single pass, so values are listed in the order they appear in the query
|
||||
# - Match on a word boundary, bind parameters can be followed by any non-word character, i.e. `in (:x)`
|
||||
params = []
|
||||
for key, value in parameters.items():
|
||||
pattern = rf"\:{key}(?=\s|$)"
|
||||
match = re.search(pattern, query)
|
||||
if match:
|
||||
query = re.sub(pattern, "?", query, count=1)
|
||||
params.append((match.start(), value))
|
||||
if parameters:
|
||||
pattern = r"\:(" + "|".join(re.escape(key) for key in parameters) + r")\b"
|
||||
params = [parameters[match.group(1)] for match in re.finditer(pattern, query)]
|
||||
query = re.sub(pattern, "?", query)
|
||||
|
||||
# Repack query and parameter list
|
||||
args = (query, [value for _, value in sorted(params, key=lambda x: x[0])])
|
||||
args = (query, params)
|
||||
|
||||
return args
|
||||
|
||||
|
||||
@@ -31,6 +31,9 @@ class Embedded(RDBMS):
|
||||
self.path = path
|
||||
|
||||
def save(self, path):
|
||||
# Initialize connection if not open
|
||||
self.initialize()
|
||||
|
||||
# Temporary database
|
||||
if not self.path:
|
||||
# Save temporary database
|
||||
|
||||
@@ -30,7 +30,7 @@ class EncoderFactory:
|
||||
# Get parent package
|
||||
encoder = ".".join(__name__.split(".")[:-1]) + "." + encoder.capitalize() + "Encoder"
|
||||
|
||||
return Resolver()(encoder)
|
||||
return Resolver()(encoder, Encoder)
|
||||
|
||||
@staticmethod
|
||||
def create(encoder):
|
||||
|
||||
@@ -6,6 +6,7 @@ from urllib.parse import urlparse
|
||||
|
||||
from ..util import Resolver
|
||||
|
||||
from .base import Database
|
||||
from .client import Client
|
||||
from .duckdb import DuckDB
|
||||
from .sqlite import SQLite
|
||||
@@ -72,6 +73,6 @@ class DatabaseFactory:
|
||||
"""
|
||||
|
||||
try:
|
||||
return Resolver()(backend)(config)
|
||||
return Resolver()(backend, Database)(config)
|
||||
except Exception as e:
|
||||
raise ImportError(f"Unable to resolve database backend: '{backend}'") from e
|
||||
|
||||
@@ -9,7 +9,6 @@ from .base import Database
|
||||
from .schema import Statement
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class RDBMS(Database):
|
||||
"""
|
||||
Base relational database class. A relational database uses SQL to insert, update, delete and select from a
|
||||
@@ -186,8 +185,9 @@ class RDBMS(Database):
|
||||
if indexids:
|
||||
select = f"{self.resolve('indexid')}, {self.resolve('score')}"
|
||||
|
||||
# Use JOIN when documents table is used to utilize indexes, default to LEFT JOIN
|
||||
join = "JOIN" if any(x and self.jsonprefix() in x for x in [where, groupby, orderby]) else "LEFT JOIN"
|
||||
# Use JOIN when a JSON field is filtered, since that filter already excludes rows without a
|
||||
# documents row. GROUP BY/ORDER BY don't filter, so they keep the LEFT JOIN.
|
||||
join = "JOIN" if where and self.jsonprefix() in where else "LEFT JOIN"
|
||||
|
||||
# Build query text
|
||||
query = Statement.TABLE_CLAUSE % (select, join)
|
||||
|
||||
@@ -3,6 +3,7 @@ SQL imports
|
||||
"""
|
||||
|
||||
from .aggregate import Aggregate
|
||||
from .base import SQL, SQLError
|
||||
from .base import SQL
|
||||
from .error import SQLError
|
||||
from .expression import Expression
|
||||
from .token import Token
|
||||
|
||||
@@ -6,6 +6,7 @@ import itertools
|
||||
import operator
|
||||
|
||||
from .base import SQL
|
||||
from .error import SQLError
|
||||
|
||||
|
||||
class Aggregate(SQL):
|
||||
@@ -32,8 +33,9 @@ class Aggregate(SQL):
|
||||
# Parse query
|
||||
query = super().__call__(query)
|
||||
|
||||
# Check if this is a SQL query
|
||||
if "select" in query:
|
||||
# Check if this is a SQL query with results. A sharded query that matches
|
||||
# nothing across all shards yields empty results, so guard before indexing.
|
||||
if "select" in query and results:
|
||||
# Get list of unique and aggregate columns. If no aggregate columns or order by found, skip
|
||||
columns = list(results[0].keys())
|
||||
aggcolumns = self.aggcolumns(columns)
|
||||
@@ -61,18 +63,41 @@ class Aggregate(SQL):
|
||||
|
||||
aggregates = {}
|
||||
for column in columns:
|
||||
column = column.lower()
|
||||
if column.startswith(("count(", "sum(", "total(")):
|
||||
# Match function names case-insensitively but keep the column name as returned
|
||||
name = column.lower()
|
||||
if name.startswith(("count(", "sum(", "total(")):
|
||||
aggregates[column] = sum
|
||||
elif column.startswith("max("):
|
||||
elif name.startswith("max("):
|
||||
aggregates[column] = max
|
||||
elif column.startswith("min("):
|
||||
elif name.startswith("min("):
|
||||
aggregates[column] = min
|
||||
elif column.startswith("avg("):
|
||||
aggregates[column] = lambda x: sum(x) / len(x)
|
||||
elif name.startswith("avg("):
|
||||
aggregates[column] = self.avg
|
||||
|
||||
return aggregates
|
||||
|
||||
def avg(self, values, counts):
|
||||
"""
|
||||
Combines per-shard/per-group average values into a single average, weighted by a
|
||||
matching count(*) column when one was selected. Without it, there is no way to
|
||||
recombine the averages correctly, so this raises instead of returning a wrong number.
|
||||
|
||||
Args:
|
||||
values: list of per-shard/per-group average values
|
||||
counts: list of matching count(*) values, or None
|
||||
|
||||
Returns:
|
||||
combined average
|
||||
"""
|
||||
|
||||
if len(values) == 1:
|
||||
return values[0]
|
||||
|
||||
if not counts or sum(counts) == 0:
|
||||
raise SQLError("avg() requires a count(*) column to combine results from multiple shards")
|
||||
|
||||
return sum(value * count for value, count in zip(values, counts)) / sum(counts)
|
||||
|
||||
def aggregate(self, query, results, columns, aggcolumns):
|
||||
"""
|
||||
Merges aggregate columns in results.
|
||||
@@ -93,6 +118,9 @@ class Aggregate(SQL):
|
||||
else:
|
||||
results = [results]
|
||||
|
||||
# Column providing the row count each group's values were computed over, used to weight avg() columns
|
||||
countcolumn = next((column for column in columns if column.lower() == "count(*)"), None)
|
||||
|
||||
# Compute column values
|
||||
rows = []
|
||||
for result in results:
|
||||
@@ -100,9 +128,10 @@ class Aggregate(SQL):
|
||||
row = {}
|
||||
for column in columns:
|
||||
if column in aggcolumns:
|
||||
# Calculate aggregate value
|
||||
function = aggcolumns[column]
|
||||
row[column] = function([r[column] for r in result])
|
||||
# Calculate aggregate value, skip NULLs from shards with no matching rows
|
||||
matches = [r for r in result if r[column] is not None]
|
||||
values, counts = [r[column] for r in matches], [r[countcolumn] for r in matches] if countcolumn else None
|
||||
row[column] = (self.avg(values, counts) if column.lower().startswith("avg(") else aggcolumns[column](values)) if values else None
|
||||
else:
|
||||
# Non aggregate column value repeat, use first value
|
||||
row[column] = result[0][column]
|
||||
@@ -154,9 +183,9 @@ class Aggregate(SQL):
|
||||
clause = clause.rsplit(" ")[0]
|
||||
reverse = True
|
||||
|
||||
# Order by columns must be in select clause
|
||||
# Order by columns must be in select clause. NULLs sort first like in SQLite (last when descending).
|
||||
if clause in query["select"]:
|
||||
results = sorted(results, key=operator.itemgetter(clause), reverse=reverse)
|
||||
results = sorted(results, key=lambda row, column=clause: (row[column] is not None, row[column]), reverse=reverse)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ SQL module
|
||||
from io import StringIO
|
||||
from shlex import shlex
|
||||
|
||||
from .error import SQLError
|
||||
from .expression import Expression
|
||||
|
||||
|
||||
@@ -42,8 +43,11 @@ class SQL:
|
||||
|
||||
clauses = None
|
||||
if self.issql(query):
|
||||
# Ignore multiple statements
|
||||
query = query.split(";")[0]
|
||||
# Multiple statements are not supported
|
||||
components = [x for x in query.split(";") if x.strip()]
|
||||
if len(components) > 1:
|
||||
raise SQLError("Multiple SQL statements are not supported. Pass values that contain a semicolon as a bind parameter.")
|
||||
query = components[0]
|
||||
|
||||
# Tokenize query
|
||||
tokens, positions = self.tokenize(query)
|
||||
@@ -66,6 +70,10 @@ class SQL:
|
||||
if similar:
|
||||
clauses["similar"] = similar
|
||||
|
||||
# Add select column aliases, if any. Only a query with an offset uses them.
|
||||
if aliases and clauses["offset"]:
|
||||
clauses["aliases"] = aliases
|
||||
|
||||
# Return clauses, default to full query if this is not a SQL query
|
||||
return clauses if clauses else {"similar": [[query]]}
|
||||
|
||||
@@ -181,9 +189,3 @@ class SQL:
|
||||
clause = self.expression(clause, alias, aliases, similar)
|
||||
|
||||
return clause
|
||||
|
||||
|
||||
class SQLError(Exception):
|
||||
"""
|
||||
Raised for errors generated by user SQL queries
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
Error module
|
||||
"""
|
||||
|
||||
|
||||
class SQLError(Exception):
|
||||
"""
|
||||
Raised for errors generated by user SQL queries
|
||||
"""
|
||||
@@ -2,6 +2,7 @@
|
||||
Expression module
|
||||
"""
|
||||
|
||||
from .error import SQLError
|
||||
from .token import Token
|
||||
|
||||
|
||||
@@ -185,6 +186,10 @@ class Expression:
|
||||
if x not in aliases.values():
|
||||
alias = tokens[x]
|
||||
|
||||
# An empty expression (e.g. a stray comma in the select clause) has no alias
|
||||
if not alias:
|
||||
raise SQLError("Empty expression in SQL clause")
|
||||
|
||||
# Strip leading/trailing brackets from alias name that doesn't have operators
|
||||
if not any(Token.isoperator(t) for t in alias) and alias[0] in ("[", "(") and alias[-1] in ("]", ")"):
|
||||
alias = alias[1:-1]
|
||||
@@ -224,7 +229,7 @@ class Expression:
|
||||
|
||||
# Read until token is a end bracket
|
||||
while token and (token != "]" or brackets > 0):
|
||||
x, token = next(iterator, (None, None))
|
||||
x, token = self.readtoken(iterator)
|
||||
|
||||
# Increase/decrease bracket counter
|
||||
if token == "[":
|
||||
@@ -263,7 +268,7 @@ class Expression:
|
||||
|
||||
# Read until token is a closing paren
|
||||
while token and token != ")":
|
||||
x, token = next(iterator, (None, None))
|
||||
x, token = self.readtoken(iterator)
|
||||
if token and token not in ["(", ",", ")"]:
|
||||
# Strip quotes and accumulate tokens
|
||||
params.append(token.replace("'", "").replace('"', ""))
|
||||
@@ -291,7 +296,7 @@ class Expression:
|
||||
|
||||
# Consume function parameters
|
||||
while token and token != ")":
|
||||
x, token = next(iterator, (None, None))
|
||||
x, token = self.readtoken(iterator)
|
||||
|
||||
# Check if token is a square bracket
|
||||
if Token.isbracket(token):
|
||||
@@ -336,6 +341,10 @@ class Expression:
|
||||
if token in Token.ALIAS:
|
||||
x, token = next(iterator, (None, None))
|
||||
|
||||
# A trailing AS with no alias name exhausts the stream
|
||||
if x is None:
|
||||
raise SQLError("Missing alias name in SQL expression")
|
||||
|
||||
# Consume tokens until end of stream or a separator is found. Evaluate next token to prevent consuming here.
|
||||
while x + 1 < len(tokens) and not Token.isseparator(Token.get(tokens, x + 1)):
|
||||
x, token = next(iterator, (None, None))
|
||||
@@ -402,3 +411,20 @@ class Expression:
|
||||
return token
|
||||
|
||||
return self.resolver(token)
|
||||
|
||||
def readtoken(self, iterator):
|
||||
"""
|
||||
Reads the next token from the token iterator. Raises a SQLError if the iterator has no more tokens.
|
||||
|
||||
Args:
|
||||
iterator: token iterator
|
||||
|
||||
Returns:
|
||||
(x, token) for the next token
|
||||
"""
|
||||
|
||||
x, token = next(iterator, (None, None))
|
||||
if x is None:
|
||||
raise SQLError("Unterminated clause in SQL expression")
|
||||
|
||||
return x, token
|
||||
|
||||
@@ -26,6 +26,9 @@ class Token:
|
||||
# Default list of sort order operators
|
||||
SORT_ORDER = ["asc", "desc"]
|
||||
|
||||
# Default list of boolean literals
|
||||
BOOLEANS = ["true", "false"]
|
||||
|
||||
@staticmethod
|
||||
def get(tokens, x):
|
||||
"""
|
||||
@@ -197,8 +200,8 @@ class Token:
|
||||
True if this token is a literal, False otherwise
|
||||
"""
|
||||
|
||||
# Literals are wrapped in quotes, parens, wildcards or numeric.
|
||||
return token and (token.startswith(("'", '"', ",", "(", ")", "*")) or token.replace(".", "", 1).isdigit())
|
||||
# Literals are wrapped in quotes, parens, wildcards, numeric or boolean.
|
||||
return token and (token.startswith(("'", '"', ",", "(", ")", "*")) or token.replace(".", "", 1).isdigit() or token.lower() in Token.BOOLEANS)
|
||||
|
||||
@staticmethod
|
||||
def islogicseparator(token):
|
||||
|
||||
@@ -18,7 +18,6 @@ from .index import Action, Configuration, Functions, Indexes, IndexIds, Reducer,
|
||||
from .search import Explain, Ids, Query, Search, Terms
|
||||
|
||||
|
||||
# pylint: disable=C0302,R0904
|
||||
class Embeddings:
|
||||
"""
|
||||
Embeddings databases are the engine that delivers semantic search. Data is transformed into embeddings vectors where similar concepts
|
||||
@@ -136,7 +135,7 @@ class Embeddings:
|
||||
self.ann.index(embeddings)
|
||||
|
||||
# Save indexids-ids mapping for indexes with no database, except when this is a reindex
|
||||
if ids and not reindex and not self.database:
|
||||
if ids is not None and not reindex and not self.database:
|
||||
self.ids = self.createids(ids)
|
||||
|
||||
# Index scoring, if necessary
|
||||
@@ -488,7 +487,8 @@ class Embeddings:
|
||||
query reduced down to keyword terms
|
||||
"""
|
||||
|
||||
return self.batchterms([query])[0]
|
||||
results = self.batchterms([query])
|
||||
return results[0] if results else None
|
||||
|
||||
def batchterms(self, queries):
|
||||
"""
|
||||
@@ -548,6 +548,9 @@ class Embeddings:
|
||||
if cloud:
|
||||
path = cloud.load(path)
|
||||
|
||||
# Reset archive since this replaces the current index
|
||||
self.archive = None
|
||||
|
||||
# Check if this is an archive file and extract
|
||||
path, apath = self.checkarchive(path)
|
||||
if apath:
|
||||
@@ -930,7 +933,8 @@ class Embeddings:
|
||||
"""
|
||||
|
||||
# Create archive instance, if necessary
|
||||
self.archive = ArchiveFactory.create()
|
||||
if not self.archive:
|
||||
self.archive = ArchiveFactory.create()
|
||||
|
||||
# Check if path is an archive file
|
||||
if self.archive.isarchive(path):
|
||||
|
||||
@@ -147,10 +147,10 @@ class Transform:
|
||||
if not self.indexing and not document[1].get(self.text):
|
||||
document[1][self.text] = str(document[0])
|
||||
|
||||
if self.text in document[1]:
|
||||
if document[1].get(self.text) is not None:
|
||||
yield (document[0], document[1][self.text], document[2])
|
||||
offset += 1
|
||||
elif self.object in document[1]:
|
||||
elif document[1].get(self.object) is not None:
|
||||
yield (document[0], document[1][self.object], document[2])
|
||||
offset += 1
|
||||
else:
|
||||
|
||||
@@ -92,7 +92,7 @@ class Search:
|
||||
Args:
|
||||
queries: list of queries
|
||||
limit: maximum results
|
||||
weights: hybrid score weights
|
||||
weights: hybrid score weights, either one set of weights for all queries or a list of [dense, sparse] weights per query
|
||||
index: index name
|
||||
|
||||
Returns:
|
||||
@@ -114,9 +114,13 @@ class Search:
|
||||
if isinstance(weights, (int, float)):
|
||||
weights = [weights, 1 - weights]
|
||||
|
||||
# Use the same weights for each query, unless a list of weights per query is passed
|
||||
if not isinstance(weights[0], (list, tuple)):
|
||||
weights = [weights] * len(queries)
|
||||
|
||||
# Create weighted scores via hybrid fusion strategy
|
||||
fusion = Hybrid(self.scoring)
|
||||
return [fusion(vectors, weights, limit) for vectors in zip(dense, sparse)]
|
||||
return [fusion(vectors, weights[x], limit) for x, vectors in enumerate(zip(dense, sparse))]
|
||||
|
||||
# Raise an error if when no indexes are available
|
||||
if not sparse and not dense:
|
||||
@@ -228,8 +232,8 @@ class Search:
|
||||
# Override limit with query limit, if applicable
|
||||
limit = max(limit, self.limit(queries))
|
||||
|
||||
# Bulk index scan
|
||||
scan = Scan(self.search, limit, weights, index)(queries, parameters)
|
||||
# Bulk index scan. The database query skips OFFSET rows, so add them to the number of candidates.
|
||||
scan = Scan(self.search, limit + self.offset(queries), weights, index)(queries, parameters)
|
||||
|
||||
# Combine index search results with database search results
|
||||
results = []
|
||||
@@ -287,15 +291,41 @@ class Search:
|
||||
# Override limit with largest limit from database queries
|
||||
qlimit = 0
|
||||
for query in queries:
|
||||
# Parse out qlimit
|
||||
# Parse out qlimit; a non-numeric limit (e.g. a ":n" bind parameter) isn't a
|
||||
# candidate-count override and must not be compared as a str against an int.
|
||||
l = query.get("limit")
|
||||
if l and l.isdigit():
|
||||
l = int(l)
|
||||
l = int(l) if l and l.isdigit() else 0
|
||||
|
||||
qlimit = l if l and l > qlimit else qlimit
|
||||
qlimit = l if l > qlimit else qlimit
|
||||
|
||||
return qlimit
|
||||
|
||||
def offset(self, queries):
|
||||
"""
|
||||
Parses the largest OFFSET clause from queries ordered by score.
|
||||
|
||||
Args:
|
||||
queries: list of queries
|
||||
|
||||
Returns:
|
||||
largest offset number or 0 if not found
|
||||
"""
|
||||
|
||||
qoffset = 0
|
||||
for query in queries:
|
||||
# Skip non-numeric offsets (e.g. a ":n" bind parameter)
|
||||
o = query.get("offset")
|
||||
o = int(o) if o and o.isdigit() else 0
|
||||
|
||||
# Only a query ordered first by score desc counts its offset. Any other order, or a select alias named score that replaces
|
||||
# the similarity score in the order, must see the same candidates on every page.
|
||||
if o > qoffset:
|
||||
orderby = query.get("orderby")
|
||||
if (not orderby or orderby.split(",")[0].strip().lower() == "score desc") and "score" not in query.get("aliases", {}):
|
||||
qoffset = o
|
||||
|
||||
return qoffset
|
||||
|
||||
def graphsearch(self, queries, limit, weights, index):
|
||||
"""
|
||||
Executes an index + graph search.
|
||||
@@ -316,8 +346,8 @@ class Search:
|
||||
# Override limit with query limit, if applicable
|
||||
limit = max(limit, self.limit(queries))
|
||||
|
||||
# Bulk index scan
|
||||
scan = Scan(self.search, limit, weights, index)(queries, None)
|
||||
# Bulk index scan. The graph query drops SKIP rows, so add them to the number of candidates.
|
||||
scan = Scan(self.search, limit + self.offset(queries), weights, index)(queries, None)
|
||||
|
||||
# Combine index search results with database search results
|
||||
for x, query in enumerate(queries):
|
||||
|
||||
@@ -2,11 +2,6 @@
|
||||
Explain module
|
||||
"""
|
||||
|
||||
# Core library imports
|
||||
from ...util import Library
|
||||
|
||||
np = Library().numpy()
|
||||
|
||||
|
||||
class Explain:
|
||||
"""
|
||||
@@ -91,8 +86,9 @@ class Explain:
|
||||
# Parse query
|
||||
query = self.database.parse(query)
|
||||
|
||||
# Extract query from similar clause
|
||||
query = " ".join([" ".join(clause) for clause in query["similar"]]) if "similar" in query else None
|
||||
# Extract query text from similar clauses. Only the first argument is query text, the rest are
|
||||
# candidates, index and weights.
|
||||
query = " ".join([clause[0] for clause in query["similar"]]) if "similar" in query else None
|
||||
|
||||
# Return original texts if query, text or score not present
|
||||
if not query or not texts or "score" not in texts[0] or "text" not in texts[0]:
|
||||
@@ -111,7 +107,7 @@ class Explain:
|
||||
permutations.append([" ".join(data)])
|
||||
|
||||
# Calculate similarity for each input text permutation and get score delta as importance
|
||||
scores = [(i, result["score"] - np.abs(s)) for i, s in self.embeddings.similarity(query, permutations)]
|
||||
scores = [(i, result["score"] - abs(s)) for i, s in self.embeddings.similarity(query, permutations)]
|
||||
|
||||
# Append tokens to result
|
||||
result["tokens"] = [(tokens[i], score) for i, score in sorted(scores, key=lambda x: x[0])]
|
||||
|
||||
@@ -194,7 +194,8 @@ class LogOdds:
|
||||
array = [s for s in raw if s > 0]
|
||||
if array:
|
||||
median = sorted(array)[len(array) // 2]
|
||||
std = (sum((x - sum(array) / len(array)) ** 2 for x in array) / len(array)) ** 0.5
|
||||
mean = sum(array) / len(array)
|
||||
std = (sum((x - mean) ** 2 for x in array) / len(array)) ** 0.5
|
||||
alpha = 1.0 / std if std > 0 else 1.0
|
||||
|
||||
return median, alpha
|
||||
|
||||
@@ -61,9 +61,9 @@ class Scan:
|
||||
|
||||
candidates = max(candidates) if candidates else default
|
||||
|
||||
# Query weights to pass to batch search
|
||||
weights = [query.weights for query in iqueries if query.weights is not None]
|
||||
weights = max(weights) if weights else self.weights
|
||||
# Query weights to pass to batch search. Each query keeps its own weights, queries without weights use the default.
|
||||
weights = [query.weights if query.weights is not None else self.weights for query in iqueries]
|
||||
weights = [[weight, 1 - weight] if isinstance(weight, (int, float)) else weight for weight in weights]
|
||||
|
||||
# Index to run query against
|
||||
index = index if index else self.index
|
||||
|
||||
@@ -37,8 +37,10 @@ class Terms:
|
||||
# Parse query
|
||||
parse = self.database.parse(query)
|
||||
|
||||
# Join terms from similar clauses
|
||||
terms.append(" ".join(" ".join(s) for s in parse["similar"]))
|
||||
# Join query text from similar clauses (absent when the query has no similar() clause).
|
||||
# Only the first argument is query text, the rest are candidates, index and weights.
|
||||
if "similar" in parse:
|
||||
terms.append(" ".join(s[0] for s in parse["similar"]))
|
||||
|
||||
return terms
|
||||
|
||||
|
||||
@@ -7,7 +7,6 @@ from collections import Counter
|
||||
from .topics import Topics
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class Graph:
|
||||
"""
|
||||
Base class for Graph instances. This class builds graph networks. Supports topic modeling
|
||||
@@ -460,6 +459,10 @@ class Graph:
|
||||
|
||||
# Also delete topic, if it's empty
|
||||
if not self.topics[topic]:
|
||||
# Categories are stored by topic position
|
||||
if self.categories:
|
||||
self.categories.pop(list(self.topics).index(topic))
|
||||
|
||||
self.topics.pop(topic)
|
||||
|
||||
# Delete node
|
||||
@@ -536,19 +539,21 @@ class Graph:
|
||||
# Unpack node and score, if available
|
||||
node, score = node if isinstance(node, tuple) else (node, None)
|
||||
|
||||
# Add nodes
|
||||
graph.addnode(node, **self.node(node))
|
||||
# Add node if it exists
|
||||
if self.hasnode(node):
|
||||
# Add nodes
|
||||
graph.addnode(node, **self.node(node))
|
||||
|
||||
# Add score if present
|
||||
if score is not None:
|
||||
graph.addattribute(node, "score", score)
|
||||
# Add score if present
|
||||
if score is not None:
|
||||
graph.addattribute(node, "score", score)
|
||||
|
||||
# Add edges
|
||||
edges = self.edges(node)
|
||||
if edges:
|
||||
for target, attributes in self.edges(node).items():
|
||||
if target in nodeids:
|
||||
graph.addedge(node, target, **attributes)
|
||||
# Add edges
|
||||
edges = self.edges(node)
|
||||
if edges:
|
||||
for target, attributes in self.edges(node).items():
|
||||
if target in nodeids:
|
||||
graph.addedge(node, target, **attributes)
|
||||
|
||||
# Filter categories and topics
|
||||
if self.topics:
|
||||
@@ -742,7 +747,7 @@ class Graph:
|
||||
def infertopics(self):
|
||||
"""
|
||||
Infers topics for all nodes with an "updated" attribute. This method analyzes the direct node
|
||||
neighbors and set the most commonly occuring topic and category for each node.
|
||||
neighbors and set the most commonly occurring topic and category for each node.
|
||||
"""
|
||||
|
||||
# Iterate over nodes missing topic attribute (only occurs for new nodes)
|
||||
@@ -755,7 +760,7 @@ class Graph:
|
||||
ids = ids.keys() if ids else None
|
||||
|
||||
# Infer topic
|
||||
topic = Counter(self.attribute(x, "topic") for x in ids).most_common(1)[0][0] if ids else None
|
||||
topic = self.commonattribute(ids, "topic")
|
||||
if topic:
|
||||
# Add id to topic list and set topic attribute
|
||||
self.topics[topic].append(node)
|
||||
@@ -765,5 +770,21 @@ class Graph:
|
||||
self.addattribute(node, "topicrank", len(self.topics[topic]) - 1)
|
||||
|
||||
# Infer category
|
||||
category = Counter(self.attribute(x, "category") for x in ids).most_common(1)[0][0]
|
||||
self.addattribute(node, "category", category)
|
||||
category = self.commonattribute([x for x in ids if self.attribute(x, "topic") == topic], "category")
|
||||
if category:
|
||||
self.addattribute(node, "category", category)
|
||||
|
||||
def commonattribute(self, ids, attribute):
|
||||
"""
|
||||
Gets the most common non-empty attribute value across a list of nodes.
|
||||
|
||||
Args:
|
||||
ids: list of node ids
|
||||
attribute: attribute name
|
||||
|
||||
Returns:
|
||||
most common value or None if no node has the attribute
|
||||
"""
|
||||
|
||||
values = Counter(value for value in (self.attribute(x, attribute) for x in ids or []) if value is not None)
|
||||
return values.most_common(1)[0][0] if values else None
|
||||
|
||||
@@ -4,6 +4,7 @@ Factory module
|
||||
|
||||
from ..util import Resolver
|
||||
|
||||
from .base import Graph
|
||||
from .networkx import NetworkX
|
||||
from .rdbms import RDBMS
|
||||
|
||||
@@ -56,6 +57,6 @@ class GraphFactory:
|
||||
"""
|
||||
|
||||
try:
|
||||
return Resolver()(backend)(config)
|
||||
return Resolver()(backend, Graph)(config)
|
||||
except Exception as e:
|
||||
raise ImportError(f"Unable to resolve graph backend: '{backend}'") from e
|
||||
|
||||
@@ -24,7 +24,6 @@ from .base import Graph
|
||||
from .query import Query
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class NetworkX(Graph):
|
||||
"""
|
||||
Graph instance backed by NetworkX.
|
||||
@@ -131,7 +130,7 @@ class NetworkX(Graph):
|
||||
for value in column:
|
||||
if isinstance(value, list):
|
||||
# Path group
|
||||
nodes.update([node for node in value if node and not isinstance(node, dict)])
|
||||
nodes.update([node for node in value if node is not None and not isinstance(node, dict)])
|
||||
elif isinstance(value, dict):
|
||||
# Nodes by id attribute
|
||||
nodes.update(uid for uid, attr in self.scan(data=True) if attr["id"] == value["id"])
|
||||
|
||||
@@ -88,7 +88,7 @@ class Query:
|
||||
"""
|
||||
|
||||
# Parameters
|
||||
where, limit, nodes, similar = None, None, [], []
|
||||
where, limit, offset, nodes, similar = None, None, None, [], []
|
||||
|
||||
# Parse where clause
|
||||
match = re.search(r"where(.+?)return", query, flags=re.DOTALL | re.IGNORECASE)
|
||||
@@ -100,6 +100,13 @@ class Query:
|
||||
if match:
|
||||
limit = match.group(1)
|
||||
|
||||
# Parse skip clause when results are ordered first by score desc. Graph results with no ORDER BY do not follow score order.
|
||||
# Most queries have no SKIP, so check for the keyword before the regex.
|
||||
if "skip" in query.lower():
|
||||
match = re.search(r"skip\s+(\d+)", query, flags=re.DOTALL | re.IGNORECASE)
|
||||
if match and re.search(r"order\s+by\s+\w+\.score\s+desc", query, flags=re.DOTALL | re.IGNORECASE):
|
||||
offset = match.group(1)
|
||||
|
||||
# Parse similar clauses
|
||||
for x, match in enumerate(re.finditer(r"similar\((.+?)\)", query, flags=re.DOTALL | re.IGNORECASE)):
|
||||
# Replace similar clause with placeholder
|
||||
@@ -115,6 +122,7 @@ class Query:
|
||||
"query": query,
|
||||
"where": where,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"nodes": nodes,
|
||||
"similar": similar,
|
||||
}
|
||||
|
||||
@@ -59,6 +59,10 @@ class RDBMS(NetworkX):
|
||||
else:
|
||||
yield from super().scan(attribute, data)
|
||||
|
||||
def addattribute(self, node, field, value):
|
||||
if self.hasnode(node):
|
||||
self.graph.backend.add_node(node, {field: value})
|
||||
|
||||
def load(self, path):
|
||||
# Create graph instance
|
||||
self.graph, self.database = self.connect()
|
||||
|
||||
@@ -82,8 +82,10 @@ class Topics:
|
||||
# Get topn terms
|
||||
terms = self.topn(idf, topn)
|
||||
|
||||
# Sort community by score descending
|
||||
community = [uid for uid, _ in scoring.search(terms, len(community))]
|
||||
# Sort community by score descending. Ensure both scored and unscored nodes are kept in community.
|
||||
ranked = [uid for uid, _ in scoring.search(terms, len(community))]
|
||||
unranked = sorted(set(community) - set(ranked), key=lambda x: centrality[x], reverse=True)
|
||||
community = ranked + unranked
|
||||
else:
|
||||
# No text found for topic, generate topic name
|
||||
terms = ["topic", str(index)]
|
||||
|
||||
@@ -7,4 +7,7 @@ from .cls import ClsPooling
|
||||
from .factory import PoolingFactory
|
||||
from .last import LastPooling
|
||||
from .late import LatePooling
|
||||
from .lemur import Lemur
|
||||
from .max import MaxPooling
|
||||
from .mean import MeanPooling
|
||||
from .muvera import Muvera
|
||||
|
||||
@@ -36,7 +36,9 @@ class Pooling(Module):
|
||||
super().__init__()
|
||||
|
||||
self.model = Models.load(path, modelargs=modelargs)
|
||||
self.tokenizer = Models.tokenizer(tokenizer if tokenizer else path)
|
||||
self.tokenizer = Models.tokenizer(
|
||||
tokenizer if tokenizer else path, trust_remote_code=modelargs.get("trust_remote_code", False) if modelargs else False
|
||||
)
|
||||
self.device = Models.device(device)
|
||||
|
||||
# Detect unbounded tokenizer typically found in older models
|
||||
|
||||
@@ -9,6 +9,7 @@ from .base import Pooling
|
||||
from .cls import ClsPooling
|
||||
from .last import LastPooling
|
||||
from .late import LatePooling
|
||||
from .max import MaxPooling
|
||||
from .mean import MeanPooling
|
||||
|
||||
from ...util import Download, DownloadError
|
||||
@@ -44,7 +45,7 @@ class PoolingFactory:
|
||||
return Pooling(path, device, tokenizer, maxlength, loadprompts, modelargs)
|
||||
|
||||
# Derive pooling method if it's not specified and path is a string
|
||||
if (not method or method not in ("clspooling", "meanpooling", "lastpooling", "latepooling")) and isinstance(path, str):
|
||||
if (not method or method not in ("clspooling", "lastpooling", "latepooling", "maxpooling", "meanpooling")) and isinstance(path, str):
|
||||
method = PoolingFactory.method(path)
|
||||
|
||||
# Check for cls pooling
|
||||
@@ -59,6 +60,10 @@ class PoolingFactory:
|
||||
if method == "latepooling":
|
||||
return LatePooling(path, device, tokenizer, maxlength, loadprompts, modelargs)
|
||||
|
||||
# Check for max pooling
|
||||
if method == "maxpooling":
|
||||
return MaxPooling(path, device, tokenizer, maxlength, loadprompts, modelargs)
|
||||
|
||||
# Default to mean pooling
|
||||
return MeanPooling(path, device, tokenizer, maxlength, loadprompts, modelargs)
|
||||
|
||||
@@ -80,14 +85,22 @@ class PoolingFactory:
|
||||
# Load 1_Pooling/config.json file
|
||||
config = PoolingFactory.load(path, "1_Pooling/config.json")
|
||||
|
||||
# Models trained with Sentence Transformers v5.4+ use `pooling_mode` parameter
|
||||
# Support both the new and legacy parameters
|
||||
mode = config.get("pooling_mode") if config else None
|
||||
|
||||
# Set to CLS pooling if it's enabled and mean pooling is disabled
|
||||
if config and config.get("pooling_mode_cls_token") and not config["pooling_mode_mean_tokens"]:
|
||||
if config and config.get("pooling_mode_cls_token", mode == "cls") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
|
||||
method = "clspooling"
|
||||
|
||||
# Set to last token pooling if it's enabled and mean pooling is disabled
|
||||
if config and config.get("pooling_mode_lasttoken") and not config["pooling_mode_mean_tokens"]:
|
||||
if config and config.get("pooling_mode_lasttoken", mode == "lasttoken") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
|
||||
method = "lastpooling"
|
||||
|
||||
# Set to max pooling if it's enabled and mean pooling is disabled
|
||||
if config and config.get("pooling_mode_max_tokens", mode == "max") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
|
||||
method = "maxpooling"
|
||||
|
||||
# Check for late interaction pooling
|
||||
if not config:
|
||||
# Load 1_Dense/config.json
|
||||
|
||||
@@ -3,6 +3,7 @@ Late module
|
||||
"""
|
||||
|
||||
from .base import Pooling
|
||||
from .lemur import Lemur
|
||||
from .muvera import Muvera
|
||||
|
||||
# Core library imports
|
||||
@@ -23,12 +24,18 @@ class LatePooling(Pooling):
|
||||
def __init__(self, path, device, tokenizer=None, maxlength=None, loadprompts=None, modelargs=None):
|
||||
# Check if fixed dimensional encoder is enabled
|
||||
modelargs = modelargs.copy() if modelargs else {}
|
||||
muvera = modelargs.pop("muvera", {})
|
||||
self.encoder = Muvera(**muvera) if muvera is not None else None
|
||||
muvera = self.encodersettings(modelargs.pop("muvera", {}), "muvera")
|
||||
lemur = self.encodersettings(modelargs.pop("lemur", None), "lemur")
|
||||
centerconfigured = "center" in modelargs
|
||||
center = modelargs.pop("center", None)
|
||||
|
||||
# Call parent initialization
|
||||
super().__init__(path, device, tokenizer, maxlength, loadprompts, modelargs)
|
||||
|
||||
# Create fixed dimensional encoder
|
||||
self.encoder = Lemur(device=self.device, **lemur) if lemur is not None else Muvera(**muvera) if muvera is not None else None
|
||||
self.lengths = None
|
||||
|
||||
# Get linear weights paths
|
||||
config = self.load(path, "modules.json")
|
||||
if config:
|
||||
@@ -44,6 +51,10 @@ class LatePooling(Pooling):
|
||||
# Load linear model
|
||||
self.linear = self.loadlinear(path, models)
|
||||
|
||||
# Configure token centering
|
||||
self.center = self.centersettings(center, self.linear, centerconfigured)
|
||||
self.batches = None
|
||||
|
||||
def forward(self, **inputs):
|
||||
"""
|
||||
Runs late pooling on token embeddings.
|
||||
@@ -55,6 +66,13 @@ class LatePooling(Pooling):
|
||||
Late pooled embeddings using output token embeddings (i.e. last hidden state)
|
||||
"""
|
||||
|
||||
# Track true token counts
|
||||
self.lengths.extend(inputs["attention_mask"].sum(dim=1).cpu().tolist())
|
||||
|
||||
# Track model batch boundaries for batch-scoped centering
|
||||
if self.center and self.center["scope"] == "batch":
|
||||
self.batches.append(inputs["attention_mask"].shape[0])
|
||||
|
||||
# Run through transformers model
|
||||
tokens = super().forward(**inputs)
|
||||
|
||||
@@ -70,6 +88,12 @@ class LatePooling(Pooling):
|
||||
category: embeddings category (query or data)
|
||||
"""
|
||||
|
||||
# Reset true token counts for this encoding pass
|
||||
self.lengths = []
|
||||
|
||||
if self.center and self.center["scope"] == "batch":
|
||||
self.batches = []
|
||||
|
||||
results = []
|
||||
|
||||
# Apply prefix
|
||||
@@ -98,25 +122,156 @@ class LatePooling(Pooling):
|
||||
normalized results with padding
|
||||
"""
|
||||
|
||||
length = 0
|
||||
for vectors in results:
|
||||
# Get max length
|
||||
if vectors.shape[0] > length:
|
||||
length = vectors.shape[0]
|
||||
data, maxlength = [], 0
|
||||
for vectors, length in zip(results, self.lengths):
|
||||
# Calculate max length
|
||||
maxlength = max(length, maxlength)
|
||||
|
||||
# Normalize vectors
|
||||
# Normalize vector
|
||||
vectors = vectors[:length]
|
||||
vectors /= np.linalg.norm(vectors, axis=1)[:, np.newaxis]
|
||||
data.append(vectors)
|
||||
|
||||
# Pad values
|
||||
data = []
|
||||
for vectors in results:
|
||||
data.append(np.pad(vectors, [(0, length - vectors.shape[0]), (0, 0)]))
|
||||
# Apply vector centering
|
||||
if self.center:
|
||||
data = self.centerdata(data)
|
||||
|
||||
# Fixed dimensional output
|
||||
if self.encoder:
|
||||
return self.encoder(data, category)
|
||||
|
||||
# Standard multi-vector output - pad values to maxlength
|
||||
padded = []
|
||||
for vectors in data:
|
||||
padded.append(np.pad(vectors, [(0, maxlength - vectors.shape[0]), (0, 0)]))
|
||||
|
||||
# Build NumPy array
|
||||
data = np.asarray(data)
|
||||
return np.asarray(padded)
|
||||
|
||||
# Apply fixed dimesional encoder, if necessary
|
||||
return self.encoder(data, category) if self.encoder else data
|
||||
@staticmethod
|
||||
def encodersettings(settings, name):
|
||||
"""
|
||||
Validates and resolves fixed encoder settings.
|
||||
|
||||
Args:
|
||||
settings: fixed encoder configuration
|
||||
name: fixed encoder name
|
||||
|
||||
Returns:
|
||||
resolved fixed encoder configuration
|
||||
"""
|
||||
|
||||
if settings is None or settings is False:
|
||||
return None
|
||||
|
||||
if isinstance(settings, dict):
|
||||
return settings
|
||||
|
||||
if name == "lemur" and isinstance(settings, str):
|
||||
return {"path": settings}
|
||||
|
||||
if name == "muvera" and settings is True:
|
||||
return {}
|
||||
|
||||
expected = "a path string or a dict of settings" if name == "lemur" else "a boolean or a dict of settings"
|
||||
raise ValueError(f"{name} expects {expected}")
|
||||
|
||||
def centersettings(self, center, linear, configured):
|
||||
"""
|
||||
Validates and resolves token centering settings.
|
||||
|
||||
Args:
|
||||
center: token centering configuration
|
||||
linear: loaded linear model
|
||||
configured: whether center was explicitly configured
|
||||
|
||||
Returns:
|
||||
resolved token centering configuration
|
||||
"""
|
||||
|
||||
# Enable batch centering by default for multi-linear models
|
||||
artifactcenter = getattr(self.encoder, "center", None)
|
||||
if not configured and artifactcenter is not None:
|
||||
center = {"scope": "collection", "mean": artifactcenter.detach().cpu().numpy()}
|
||||
elif not configured:
|
||||
center = sum(isinstance(module, torch.nn.Linear) for module in linear.modules()) > 1
|
||||
|
||||
if isinstance(center, bool):
|
||||
return {"scope": "batch"} if center else None
|
||||
|
||||
if not isinstance(center, dict):
|
||||
raise ValueError("center must be a boolean or dictionary")
|
||||
|
||||
center = center.copy()
|
||||
scope = center.pop("scope", "document")
|
||||
if scope not in ("document", "batch", "collection"):
|
||||
raise ValueError("center scope must be one of: document, batch, collection")
|
||||
|
||||
mean = center.pop("mean", None)
|
||||
path = center.pop("path", None)
|
||||
if center:
|
||||
raise ValueError(f"unknown center setting: {next(iter(center))}")
|
||||
|
||||
if scope != "collection" and (mean is not None or path is not None):
|
||||
raise ValueError("center mean and path are only valid with collection scope")
|
||||
|
||||
if scope == "collection":
|
||||
if (mean is None) == (path is None):
|
||||
raise ValueError("collection center scope requires exactly one of mean or path")
|
||||
|
||||
if path is not None:
|
||||
with safetensors.safe_open(path, framework="np") as source:
|
||||
mean = source.get_tensor("center.mean")
|
||||
else:
|
||||
mean = np.asarray(mean)
|
||||
|
||||
if mean.ndim != 1 or not np.isfinite(mean).all():
|
||||
raise ValueError("collection center mean must be a finite one-dimensional array")
|
||||
|
||||
return {"scope": scope, "mean": mean} if scope == "collection" else {"scope": scope}
|
||||
|
||||
def centerdata(self, data):
|
||||
"""
|
||||
Centers and re-normalizes true token rows.
|
||||
|
||||
Args:
|
||||
data: normalized token vectors, with optional zero padding
|
||||
|
||||
Returns:
|
||||
centered token vectors with zero padding preserved
|
||||
"""
|
||||
|
||||
array = isinstance(data, np.ndarray)
|
||||
vectors = [np.array(value, copy=True) for value in data]
|
||||
masks = [np.any(value != 0, axis=1) for value in vectors]
|
||||
scope = self.center["scope"]
|
||||
|
||||
means = [None] * len(vectors)
|
||||
if scope == "batch":
|
||||
batches = getattr(self, "batches", None)
|
||||
sizes = batches if batches and sum(batches) == len(vectors) else [len(vectors)]
|
||||
offset = 0
|
||||
for size in sizes:
|
||||
rows = [vectors[x][masks[x]] for x in range(offset, offset + size) if masks[x].any()]
|
||||
mean = np.concatenate(rows).mean(axis=0) if rows else None
|
||||
means[offset : offset + size] = [mean] * size
|
||||
offset += size
|
||||
elif scope == "collection":
|
||||
means = [self.center["mean"]] * len(vectors)
|
||||
|
||||
for x, (value, mask) in enumerate(zip(vectors, masks)):
|
||||
if mask.any():
|
||||
current = value[mask]
|
||||
average = current.mean(axis=0) if scope == "document" else means[x]
|
||||
if average is not None and average.shape != (current.shape[1],):
|
||||
raise ValueError("center mean dimension must match token vector dimension")
|
||||
|
||||
current = current - average
|
||||
norms = np.linalg.norm(current, axis=1, keepdims=True)
|
||||
current = np.divide(current, norms, out=np.zeros_like(current), where=norms != 0)
|
||||
vectors[x][mask] = current
|
||||
|
||||
return np.asarray(vectors) if array else vectors
|
||||
|
||||
def settings(self, path, config):
|
||||
"""
|
||||
@@ -136,7 +291,8 @@ class LatePooling(Pooling):
|
||||
config = self.load(path, "artifact.metadata")
|
||||
params = ["query_token_id", "query_maxlen", "doc_token_id", "doc_maxlen"]
|
||||
|
||||
return [config.get(p) for p in params]
|
||||
# Use defaults when the settings file is missing
|
||||
return [config.get(p) for p in params] if config else [None] * len(params)
|
||||
|
||||
def loadlinear(self, path, models):
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,412 @@
|
||||
"""
|
||||
LEMUR module
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
|
||||
# Core library imports
|
||||
from ...util import Download, Library
|
||||
|
||||
library = Library()
|
||||
np = library.numpy()
|
||||
safetensors = library.safetensors()
|
||||
torch = library.torch()
|
||||
Module = library.module()
|
||||
|
||||
|
||||
class Lemur:
|
||||
"""
|
||||
Implements the LEMUR (Learned Multi-Vector Retrieval) algorithm. This reduces
|
||||
late interaction multi-vector outputs to a single fixed vector.
|
||||
|
||||
LEMUR artifacts contain a feature encoder, output normalization statistics and
|
||||
a sample of learn-token vectors used to solve document weights. Query vectors
|
||||
are summed feature encodings while document vectors are ordinary least squares
|
||||
weights over the stored sample.
|
||||
|
||||
This code is based on the following:
|
||||
- Paper: https://arxiv.org/abs/2601.21853
|
||||
- GitHub: https://github.com/ejaasaari/lemur
|
||||
"""
|
||||
|
||||
def __init__(self, path=None, device=None):
|
||||
"""
|
||||
Creates a LEMUR encoder.
|
||||
|
||||
Args:
|
||||
path: local artifact directory or Hugging Face Hub path
|
||||
device: tensor device
|
||||
"""
|
||||
|
||||
self.device = torch.device(device) if device is not None else torch.device("cpu")
|
||||
self.model = None
|
||||
self.sample = None
|
||||
self.mean = None
|
||||
self.std = None
|
||||
self.center = None
|
||||
self.config = None
|
||||
self.pinv = None
|
||||
self.selectedepoch = None
|
||||
self.selectedloss = None
|
||||
self.selectionmetric = None
|
||||
|
||||
if path:
|
||||
self.load(path)
|
||||
|
||||
def __call__(self, data, category):
|
||||
"""
|
||||
Transforms multi-vector collections into fixed dimensional vectors.
|
||||
|
||||
Args:
|
||||
data: list or array of unpadded multi-vector collections
|
||||
category: embeddings category (query or data)
|
||||
|
||||
Returns:
|
||||
fixed dimensional vectors
|
||||
"""
|
||||
|
||||
if category == "query":
|
||||
vectors = self.computefeatures(data)
|
||||
elif category == "data":
|
||||
vectors = self.computeweights(data)
|
||||
else:
|
||||
raise ValueError("category must be query or data")
|
||||
|
||||
return vectors.detach().cpu().to(torch.float32).numpy()
|
||||
|
||||
def computefeatures(self, data):
|
||||
"""
|
||||
Computes summed query features.
|
||||
|
||||
Args:
|
||||
data: list or array of unpadded query token-vector collections
|
||||
|
||||
Returns:
|
||||
fixed dimensional query features
|
||||
"""
|
||||
|
||||
self.check()
|
||||
documents = self.documents(data)
|
||||
self.model.eval()
|
||||
|
||||
with torch.inference_mode():
|
||||
return torch.stack([self.model.features(document).sum(dim=0) / self.config["queryscale"] for document in documents])
|
||||
|
||||
def computeweights(self, data):
|
||||
"""
|
||||
Computes ordinary least squares document weights.
|
||||
|
||||
Args:
|
||||
data: list or array of unpadded document token-vector collections
|
||||
|
||||
Returns:
|
||||
fixed dimensional document weights
|
||||
"""
|
||||
|
||||
self.check()
|
||||
documents = self.documents(data)
|
||||
|
||||
# Precompute the feature sample SVD once per loaded encoder
|
||||
if self.pinv is None:
|
||||
self.model.eval()
|
||||
with torch.inference_mode():
|
||||
features = self.model.features(self.sample)
|
||||
left, values, right = torch.linalg.svd(features, full_matrices=False)
|
||||
tolerance = torch.finfo(values.dtype).eps * max(features.shape) * values.max()
|
||||
scales = torch.where(values > tolerance, values.reciprocal(), torch.zeros_like(values))
|
||||
self.pinv = (right.T * scales) @ left.T
|
||||
|
||||
targets = (self.maxsim(documents, self.sample) - self.mean) / self.std
|
||||
return (self.pinv @ targets).T
|
||||
|
||||
def documents(self, data):
|
||||
"""
|
||||
Converts input data to a list of float32 tensors.
|
||||
|
||||
Args:
|
||||
data: list or array of multi-vector collections
|
||||
|
||||
Returns:
|
||||
list of tensors
|
||||
"""
|
||||
|
||||
if isinstance(data, np.ndarray):
|
||||
if data.ndim == 2:
|
||||
data = [data]
|
||||
elif data.ndim == 3:
|
||||
data = list(data)
|
||||
|
||||
documents = []
|
||||
for document in data:
|
||||
if isinstance(document, np.ndarray):
|
||||
document = torch.from_numpy(document)
|
||||
elif torch.is_tensor(document):
|
||||
document = document.detach()
|
||||
else:
|
||||
document = torch.tensor(document)
|
||||
|
||||
if document.ndim != 2 or document.shape[0] == 0:
|
||||
raise ValueError("each document must be a non-empty 2D token-vector array")
|
||||
|
||||
documents.append(document.to(device=self.device, dtype=torch.float32))
|
||||
|
||||
return documents
|
||||
|
||||
def maxsim(self, documents, queries):
|
||||
"""
|
||||
Computes single-token maxsim targets.
|
||||
|
||||
Args:
|
||||
documents: list of document token-vector tensors
|
||||
queries: token vectors used as independent single-token queries
|
||||
|
||||
Returns:
|
||||
matrix with a row per query token and column per document
|
||||
"""
|
||||
|
||||
return torch.stack([(queries @ document.T).max(dim=1).values for document in documents], dim=1)
|
||||
|
||||
def save(self, path):
|
||||
"""
|
||||
Saves LEMUR artifacts.
|
||||
|
||||
Args:
|
||||
path: output artifact directory
|
||||
"""
|
||||
|
||||
self.check()
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
with open(os.path.join(path, "config.json"), "w", encoding="utf-8") as output:
|
||||
json.dump(self.config, output, indent=2, sort_keys=True)
|
||||
output.write("\n")
|
||||
|
||||
tensors = {name: value.detach().cpu().contiguous() for name, value in self.model.state_dict().items() if not name.startswith("outputlayer.")}
|
||||
tensors.update(
|
||||
{
|
||||
"lemur.mean": self.mean.detach().cpu().reshape(1),
|
||||
"lemur.std": self.std.detach().cpu().reshape(1),
|
||||
"lemur.sample": self.sample.detach().cpu().contiguous(),
|
||||
}
|
||||
)
|
||||
if self.center is not None:
|
||||
tensors["lemur.center"] = self.center.detach().cpu().contiguous()
|
||||
safetensors.torch.save_file(tensors, os.path.join(path, "model.safetensors"))
|
||||
|
||||
def load(self, path):
|
||||
"""
|
||||
Loads LEMUR artifacts from a local directory or Hugging Face Hub path.
|
||||
|
||||
Args:
|
||||
path: local artifact directory or Hugging Face Hub path
|
||||
"""
|
||||
|
||||
config = Download()(path, "config.json")
|
||||
weights = Download()(path, "model.safetensors")
|
||||
|
||||
with open(config, encoding="utf-8") as source:
|
||||
self.config = json.load(source)
|
||||
|
||||
# Encoder repos also have config.json and model.safetensors, check this is a LEMUR artifact
|
||||
if "modeltype" not in self.config:
|
||||
raise ValueError(f"{path} is not a LEMUR artifact, create one with LemurTrainer")
|
||||
|
||||
self.model = LemurModel(**self.config).to(self.device)
|
||||
training = self.config.get("training", {})
|
||||
self.selectedepoch = training.get("selectedepoch")
|
||||
self.selectedloss = training.get("selectedloss")
|
||||
self.selectionmetric = training.get("selectionmetric")
|
||||
with safetensors.safe_open(weights, framework="pt", device=str(self.device)) as source:
|
||||
tensors = {name: source.get_tensor(name) for name in source.keys()}
|
||||
|
||||
self.mean = tensors.pop("lemur.mean")[0]
|
||||
self.std = tensors.pop("lemur.std")[0]
|
||||
self.sample = tensors.pop("lemur.sample")
|
||||
self.center = tensors.pop("lemur.center", None)
|
||||
self.model.load_state_dict(tensors)
|
||||
self.model.eval()
|
||||
self.pinv = None
|
||||
|
||||
return self
|
||||
|
||||
def check(self):
|
||||
"""
|
||||
Validates that fitted or loaded LEMUR state is available.
|
||||
"""
|
||||
|
||||
if self.model is None or self.sample is None or self.mean is None or self.std is None:
|
||||
raise ValueError("LEMUR must be fitted or loaded before encoding")
|
||||
|
||||
|
||||
class LemurModel(Module):
|
||||
"""
|
||||
LEMUR feature extractor with an optional training-only linear readout.
|
||||
"""
|
||||
|
||||
# pylint: disable=R0913,R0917
|
||||
def __init__(
|
||||
self,
|
||||
modeltype,
|
||||
inputdim,
|
||||
hiddendim,
|
||||
finalhiddendim,
|
||||
layers,
|
||||
activation,
|
||||
queryscale,
|
||||
seed,
|
||||
outputdim=None,
|
||||
training=None,
|
||||
):
|
||||
"""
|
||||
Creates a LEMUR feature model.
|
||||
|
||||
Args:
|
||||
modeltype: elm or mlp
|
||||
inputdim: token-vector dimension
|
||||
hiddendim: MLP hidden dimension
|
||||
finalhiddendim: fixed output dimension
|
||||
layers: number of MLP feature layers
|
||||
activation: feature activation
|
||||
queryscale: query feature sum scale
|
||||
seed: artifact seed
|
||||
outputdim: optional number of training target documents
|
||||
training: optional fit-selection metadata
|
||||
"""
|
||||
|
||||
# Unused by model construction but retained in the serialized configuration
|
||||
del queryscale, seed, training
|
||||
super().__init__()
|
||||
|
||||
if modeltype == "elm":
|
||||
features = [
|
||||
RandomFeatures(inputdim, finalhiddendim, activation),
|
||||
torch.nn.LayerNorm(finalhiddendim, elementwise_affine=False),
|
||||
]
|
||||
elif modeltype == "mlp":
|
||||
features = []
|
||||
dimensions = [inputdim] + [hiddendim] * (layers - 1)
|
||||
for index, dimension in enumerate(dimensions):
|
||||
output = finalhiddendim if index == len(dimensions) - 1 else hiddendim
|
||||
features.extend([torch.nn.Linear(dimension, output), torch.nn.LayerNorm(output), Activation.module(activation)])
|
||||
else:
|
||||
raise ValueError("modeltype must be elm or mlp")
|
||||
|
||||
self.featureextractor = torch.nn.Sequential(*features)
|
||||
self.outputlayer = torch.nn.Linear(finalhiddendim, outputdim, bias=False) if outputdim is not None else None
|
||||
|
||||
def forward(self, data):
|
||||
"""
|
||||
Runs feature extraction and the training readout.
|
||||
|
||||
Args:
|
||||
data: input token vectors
|
||||
|
||||
Returns:
|
||||
predicted standardized maxsim targets
|
||||
"""
|
||||
|
||||
if self.outputlayer is None:
|
||||
raise ValueError("LEMUR training readout is not available in a loaded inference artifact")
|
||||
|
||||
return self.outputlayer(self.features(data))
|
||||
|
||||
def features(self, data):
|
||||
"""
|
||||
Extracts fixed dimensional token features.
|
||||
|
||||
Args:
|
||||
data: input token vectors
|
||||
|
||||
Returns:
|
||||
token features
|
||||
"""
|
||||
|
||||
return self.featureextractor(data)
|
||||
|
||||
|
||||
class RandomFeatures(Module):
|
||||
"""
|
||||
Seeded random activation features used by the ELM model.
|
||||
"""
|
||||
|
||||
def __init__(self, inputdim, outputdim, activation):
|
||||
"""
|
||||
Creates random activation features.
|
||||
|
||||
Args:
|
||||
inputdim: input token-vector dimension
|
||||
outputdim: fixed feature dimension
|
||||
activation: feature activation
|
||||
"""
|
||||
|
||||
super().__init__()
|
||||
self.register_buffer("weight", torch.randn(inputdim, outputdim))
|
||||
self.activation = activation
|
||||
self.scale = math.sqrt(2.0 / outputdim)
|
||||
|
||||
def forward(self, data):
|
||||
"""
|
||||
Projects and activates input vectors.
|
||||
|
||||
Args:
|
||||
data: input token vectors
|
||||
|
||||
Returns:
|
||||
random activation features
|
||||
"""
|
||||
|
||||
return self.scale * Activation.function(self.activation)(data @ self.weight)
|
||||
|
||||
|
||||
class Activation:
|
||||
"""
|
||||
Activation helpers.
|
||||
"""
|
||||
|
||||
# Activation name to torch.nn module name. Modules and functions are resolved on demand to keep
|
||||
# this module importable when Torch isn't installed.
|
||||
MODULES = {"relu": "ReLU", "gelu": "GELU", "silu": "SiLU", "mish": "Mish"}
|
||||
|
||||
@staticmethod
|
||||
def module(name):
|
||||
"""
|
||||
Creates an activation module.
|
||||
|
||||
Args:
|
||||
name: activation name
|
||||
|
||||
Returns:
|
||||
activation module
|
||||
"""
|
||||
|
||||
Activation.validate(name)
|
||||
return getattr(torch.nn, Activation.MODULES[name])()
|
||||
|
||||
@staticmethod
|
||||
def function(name):
|
||||
"""
|
||||
Gets an activation function.
|
||||
|
||||
Args:
|
||||
name: activation name
|
||||
|
||||
Returns:
|
||||
activation function
|
||||
"""
|
||||
|
||||
Activation.validate(name)
|
||||
return getattr(torch.nn.functional, name)
|
||||
|
||||
@staticmethod
|
||||
def validate(name):
|
||||
"""
|
||||
Validates an activation name.
|
||||
|
||||
Args:
|
||||
name: activation name
|
||||
"""
|
||||
|
||||
if name not in Activation.MODULES:
|
||||
raise ValueError(f"activation must be one of: {', '.join(Activation.MODULES)}")
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Max module
|
||||
"""
|
||||
|
||||
from .base import Pooling
|
||||
|
||||
# Core library imports
|
||||
from ...util import Library
|
||||
|
||||
torch = Library().torch()
|
||||
|
||||
|
||||
class MaxPooling(Pooling):
|
||||
"""
|
||||
Builds max pooled vectors usings outputs from a transformers model.
|
||||
"""
|
||||
|
||||
def forward(self, **inputs):
|
||||
"""
|
||||
Runs max pooling on token embeddings taking the input mask into account.
|
||||
|
||||
Args:
|
||||
inputs: model inputs
|
||||
|
||||
Returns:
|
||||
max pooled embeddings using output token embeddings (i.e. last hidden state)
|
||||
"""
|
||||
|
||||
# Run through transformers model
|
||||
tokens = super().forward(**inputs)
|
||||
mask = inputs["attention_mask"]
|
||||
|
||||
# Max pooling
|
||||
# pylint: disable=E1101
|
||||
mask = mask.unsqueeze(-1).expand(tokens.size()).bool()
|
||||
return torch.max(tokens.masked_fill(~mask, torch.finfo(tokens.dtype).min), 1).values
|
||||
@@ -6,6 +6,7 @@ Muvera module
|
||||
from ...util import Library
|
||||
|
||||
np = Library().numpy()
|
||||
torch = Library().torch()
|
||||
|
||||
|
||||
class Muvera:
|
||||
@@ -70,73 +71,72 @@ class Muvera:
|
||||
partitions = 2**self.hashes
|
||||
|
||||
# Document tracking
|
||||
lengths = np.array([len(doc) for doc in data], dtype=np.int32)
|
||||
total = np.sum(lengths)
|
||||
documents = np.repeat(np.arange(length), lengths)
|
||||
lengths = torch.tensor([len(doc) for doc in data], dtype=torch.int64)
|
||||
total = lengths.sum().item()
|
||||
documents = torch.repeat_interleave(torch.arange(length), lengths)
|
||||
|
||||
# Stack all vectors
|
||||
points = np.vstack(data).astype(np.float32)
|
||||
points = torch.cat([torch.as_tensor(doc, dtype=torch.float32) for doc in data])
|
||||
|
||||
# Output vectors
|
||||
size = self.repetitions * partitions * projection
|
||||
vectors = np.zeros((length, size), dtype=np.float32)
|
||||
vectors = torch.zeros((length, size), dtype=torch.float32)
|
||||
|
||||
# Process each repetition
|
||||
for number in range(self.repetitions):
|
||||
seed = self.seed + number
|
||||
|
||||
# Calculate the simhash
|
||||
sketches = points @ self.random(dimension, self.hashes, seed)
|
||||
sketches = torch.matmul(points, self.random(dimension, self.hashes, seed))
|
||||
|
||||
# Dimensionality reduction, if necessary
|
||||
projected = points if identity else (points @ self.reducer(dimension, projection, seed))
|
||||
projected = points if identity else torch.matmul(points, self.reducer(dimension, projection, seed))
|
||||
|
||||
# Get partition indices
|
||||
bits = (sketches > 0).astype(np.uint32)
|
||||
indices = np.zeros(total, dtype=np.uint32)
|
||||
bits = (sketches > 0).to(torch.int64)
|
||||
indices = torch.zeros(total, dtype=torch.int64)
|
||||
|
||||
# Calculate vector indices
|
||||
for x in range(self.hashes):
|
||||
indices = (indices << 1) + (bits[:, x] ^ (indices & 1))
|
||||
|
||||
# Initialize storage
|
||||
fdesum = np.zeros((length * partitions * projection,), dtype=np.float32)
|
||||
counts = np.zeros((length, partitions), dtype=np.int32)
|
||||
fdesum = torch.zeros((length * partitions * projection,), dtype=torch.float32)
|
||||
counts = torch.zeros((length * partitions,), dtype=torch.int32)
|
||||
|
||||
# Count vectors per partition per document
|
||||
np.add.at(counts, (documents, indices), 1)
|
||||
flat = documents * partitions + indices
|
||||
counts.index_add_(0, flat, torch.ones(total, dtype=torch.int32))
|
||||
counts = counts.reshape(length, partitions)
|
||||
|
||||
# Aggregate vectors using flattened indexing for efficiency
|
||||
part = documents * partitions + indices
|
||||
base = part * projection
|
||||
base = flat * projection
|
||||
|
||||
for d in range(projection):
|
||||
flat = base + d
|
||||
np.add.at(fdesum, flat, projected[:, d])
|
||||
fdesum.index_add_(0, base + d, projected[:, d])
|
||||
|
||||
# Reshape for easier manipulation
|
||||
# pylint: disable=E1121
|
||||
fdesum = fdesum.reshape(length, partitions, projection)
|
||||
|
||||
# Convert sums to averages for data category
|
||||
if category == "data":
|
||||
# Safe division (avoid divide by zero)
|
||||
counts = counts[:, :, np.newaxis]
|
||||
np.divide(fdesum, counts, out=fdesum, where=counts > 0)
|
||||
counts = counts[:, :, None]
|
||||
fdesum = torch.where(counts > 0, fdesum / counts.to(torch.float32), fdesum)
|
||||
|
||||
# Save results
|
||||
start = number * partitions * projection
|
||||
vectors[:, start : start + partitions * projection] = fdesum.reshape(length, -1)
|
||||
|
||||
return vectors
|
||||
return vectors.cpu().numpy()
|
||||
|
||||
def random(self, dimension, projection, seed):
|
||||
"""
|
||||
Generates a random matrix for simhash projections.
|
||||
|
||||
Args:
|
||||
dimensions: number of dimensions for input vectors
|
||||
projections: number of projection dimensions
|
||||
dimension: number of dimensions for input vectors
|
||||
projection: number of projection dimensions
|
||||
seed: random seed
|
||||
|
||||
Returns:
|
||||
@@ -144,7 +144,7 @@ class Muvera:
|
||||
"""
|
||||
|
||||
rng = np.random.default_rng(seed)
|
||||
return rng.normal(loc=0.0, scale=1.0, size=(dimension, projection)).astype(np.float32)
|
||||
return torch.from_numpy(rng.normal(loc=0.0, scale=1.0, size=(dimension, projection)).astype(np.float32))
|
||||
|
||||
def reducer(self, dimension, projection, seed):
|
||||
"""
|
||||
@@ -152,7 +152,7 @@ class Muvera:
|
||||
|
||||
Args:
|
||||
dimension: number of input dimensions
|
||||
projected: number of dimensions to project inputs to
|
||||
projection: number of dimensions to project inputs to
|
||||
|
||||
Returns:
|
||||
Dimensionality reduced matrix
|
||||
@@ -164,4 +164,4 @@ class Muvera:
|
||||
signs = rng.choice([-1.0, 1.0], size=dimension)
|
||||
out[np.arange(dimension), indices] = signs
|
||||
|
||||
return out
|
||||
return torch.from_numpy(out)
|
||||
|
||||
@@ -169,11 +169,15 @@ class Signal:
|
||||
audio with silence removed
|
||||
"""
|
||||
|
||||
# Process in 20ms chunks
|
||||
# Empty audio has no chunks to concatenate.
|
||||
if len(audio) == 0:
|
||||
return audio
|
||||
|
||||
# Process in chunks, including the final partial chunk
|
||||
n, offset = int(rate * (20 / 1000.0) * 2), 0
|
||||
|
||||
chunks = []
|
||||
while offset + n <= len(audio):
|
||||
while offset < len(audio):
|
||||
# Calculate energy for chunk and detection result
|
||||
chunk = audio[offset : offset + n]
|
||||
energyfreq = Signal.energy(chunk, rate)
|
||||
|
||||
@@ -92,7 +92,7 @@ class TextToSpeech(Pipeline):
|
||||
|
||||
# Streaming response
|
||||
if stream:
|
||||
return self.stream(texts, speaker, encoding)
|
||||
return self.stream(texts, speaker, encoding, **kwargs)
|
||||
|
||||
# Transform text to speech
|
||||
results = [self.execute(x, speaker, encoding, **kwargs) for x in texts]
|
||||
@@ -137,7 +137,7 @@ class TextToSpeech(Pipeline):
|
||||
|
||||
return exists
|
||||
|
||||
def stream(self, texts, speaker, encoding):
|
||||
def stream(self, texts, speaker, encoding, **kwargs):
|
||||
"""
|
||||
Iterates over texts, splits into segments and yields snippets of audio.
|
||||
This method is designed to integrate with streaming LLM generation.
|
||||
@@ -146,6 +146,7 @@ class TextToSpeech(Pipeline):
|
||||
texts: list of input texts
|
||||
speaker: speaker id
|
||||
encoding: audio encoding format
|
||||
kwargs: additional keyword args
|
||||
|
||||
Returns:
|
||||
snippets of audio as NumPy arrays or audio bytes depending on encoding parameter
|
||||
@@ -157,11 +158,11 @@ class TextToSpeech(Pipeline):
|
||||
|
||||
if x == "\n" or (x.strip().endswith(".") and len([y for y in buffer if y]) > 2):
|
||||
data, buffer = "".join(buffer), []
|
||||
yield self.execute(data, speaker, encoding)
|
||||
yield self.execute(data, speaker, encoding, **kwargs)
|
||||
|
||||
if buffer:
|
||||
data = "".join(buffer)
|
||||
yield self.execute(data, speaker, encoding)
|
||||
yield self.execute(data, speaker, encoding, **kwargs)
|
||||
|
||||
def execute(self, text, speaker, encoding, **kwargs):
|
||||
"""
|
||||
|
||||
@@ -4,6 +4,7 @@ Segment imports
|
||||
|
||||
from .filetohtml import FileToHTML
|
||||
from .htmltomd import HTMLToMarkdown
|
||||
from .safeopen import SafeOpen
|
||||
from .segmentation import Segmentation
|
||||
from .tabular import Tabular
|
||||
from .textractor import Textractor
|
||||
|
||||
@@ -4,6 +4,8 @@ HTMLToMarkdown module
|
||||
|
||||
import re
|
||||
|
||||
from textwrap import indent
|
||||
|
||||
# Conditional import
|
||||
try:
|
||||
from bs4 import BeautifulSoup, NavigableString
|
||||
@@ -58,6 +60,14 @@ class HTMLToMarkdown(Pipeline):
|
||||
for script in soup.find_all(["script", "style"]):
|
||||
script.decompose()
|
||||
|
||||
# Replace line breaks with newlines. Table cells and headings must stay on a single line, use a space there.
|
||||
for br in soup.find_all("br"):
|
||||
parent = br.parent
|
||||
br.replace_with(" " if br.find_parent(["th", "td", "h1", "h2", "h3", "h4", "h5", "h6"]) else "\n")
|
||||
|
||||
# Merge the line break with the text around it
|
||||
parent.smooth()
|
||||
|
||||
# Check for article sections
|
||||
article = next((x for x in ["article", "main"] if soup.find(x)), None)
|
||||
|
||||
@@ -135,7 +145,7 @@ class HTMLToMarkdown(Pipeline):
|
||||
metadata = [f"**{title.text.strip()}**"] if title and title.text else []
|
||||
|
||||
description = node.find("meta", attrs={"name": "description"})
|
||||
if description and description["content"]:
|
||||
if description and description.get("content"):
|
||||
metadata.append(f"\n*{description['content'].strip()}*")
|
||||
|
||||
# Add separator
|
||||
@@ -184,12 +194,17 @@ class HTMLToMarkdown(Pipeline):
|
||||
target, text = x if x.name else node, x.text
|
||||
|
||||
if text.strip():
|
||||
# Leading whitespace separates this text from the previous word, keep it outside of the formatting
|
||||
space = " " if text[0].isspace() and texts and not texts[-1][-1:].isspace() else ""
|
||||
|
||||
if target.name in ("b", "strong"):
|
||||
text = f"**{text.strip()}** "
|
||||
text = f"{space}**{text.strip()}** "
|
||||
elif target.name in ("i", "em"):
|
||||
text = f"*{text.strip()}* "
|
||||
text = f"{space}*{text.strip()}* "
|
||||
elif target.name == "a":
|
||||
text = f"[{text.strip()}]({target.get('href')}) "
|
||||
# Links without a target are added as plain text
|
||||
href = target.get("href")
|
||||
text = f"{space}[{text.strip()}]({href}) " if href else f"{space}{text.strip()} "
|
||||
|
||||
texts.append(text)
|
||||
|
||||
@@ -253,21 +268,54 @@ class HTMLToMarkdown(Pipeline):
|
||||
list as markdown
|
||||
"""
|
||||
|
||||
elements = []
|
||||
for x, element in enumerate(node.find_all("li")):
|
||||
# Unordered lists use dashes. Ordered lists use numbers.
|
||||
prefix = "-" if node.name == "ul" else f"{x + 1}."
|
||||
elements, count = [], 0
|
||||
if node.name == "ol":
|
||||
try:
|
||||
count = int(node.get("start", 1)) - 1
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# List item text
|
||||
text = self.process(element, article)
|
||||
for element in node.find_all(("li", "ul", "ol"), recursive=False):
|
||||
# Nested list that is a direct child of this list
|
||||
if element.name != "li":
|
||||
self.nested(element, article, elements)
|
||||
else:
|
||||
# Unordered lists use dashes. Ordered lists use numbers.
|
||||
count += 1
|
||||
prefix = "-" if node.name == "ul" else f"{count}."
|
||||
|
||||
# Add list element
|
||||
if text:
|
||||
elements.append(f"{prefix} {text}")
|
||||
# Detach nested lists, these are formatted separately
|
||||
nested = [child.extract() for child in element.find_all(("ul", "ol"), recursive=False)]
|
||||
|
||||
# List item text
|
||||
text = self.process(element, article)
|
||||
|
||||
# Add list element
|
||||
if text:
|
||||
elements.append(f"{prefix} {text}")
|
||||
|
||||
# Add nested lists
|
||||
for child in nested:
|
||||
self.nested(child, article, elements)
|
||||
|
||||
# Join elements together as string
|
||||
return "\n".join(elements)
|
||||
|
||||
def nested(self, node, article, elements):
|
||||
"""
|
||||
Nested list handler. Formats a nested list and adds it to elements indented with a tab,
|
||||
which is preserved by text cleaning.
|
||||
|
||||
Args:
|
||||
node: nested list node
|
||||
article: True if the main section node is an article
|
||||
elements: list of formatted list elements
|
||||
"""
|
||||
|
||||
text = self.items(node, article)
|
||||
if text:
|
||||
elements.append(indent(text, "\t"))
|
||||
|
||||
def code(self, node):
|
||||
"""
|
||||
Code block handler. This method transforms a HTML pre or code block into a Markdown formatted
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
"""
|
||||
SafeOpen module
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from .urlretrieve import URLRetrieve
|
||||
|
||||
|
||||
class SafeOpen:
|
||||
"""
|
||||
Safely open a local file or URL. When `safeopen` is enabled, the following rules are applied.
|
||||
|
||||
- Local files must be in the safeopen directory (defaults to temp dir)
|
||||
- URLs must be public URLs (also has checks for redirects and DNS rebind attempts)
|
||||
"""
|
||||
|
||||
def __init__(self, headers=None, safeopen=False, allowurl=True):
|
||||
"""
|
||||
Creates a new safeopen instance.
|
||||
|
||||
Args:
|
||||
safeopen: if safe validation checks should be enabled
|
||||
allowurl: if urls should be allowed
|
||||
"""
|
||||
|
||||
# HTTP headers
|
||||
self.headers = headers if headers else {}
|
||||
|
||||
# Safe open mode. When set only local temp urls (or a specified directory) and non-private URLs are supported
|
||||
self.safeopen = os.path.realpath(tempfile.gettempdir() if isinstance(safeopen, bool) else safeopen) if safeopen else safeopen
|
||||
|
||||
# URL retriever
|
||||
self.urlretrieve = URLRetrieve(self.headers, self.safeopen)
|
||||
|
||||
# Enable URL validation
|
||||
self.allowurl = allowurl
|
||||
|
||||
def valid(self, path):
|
||||
"""
|
||||
Checks if path is a valid local file or web url. Returns path if valid along with a flag
|
||||
denoting if the path exists locally.
|
||||
|
||||
Args:
|
||||
path: path to check
|
||||
|
||||
Returns:
|
||||
(path, exists)
|
||||
"""
|
||||
|
||||
# Convert file urls to local paths
|
||||
path = path.replace("file://", "")
|
||||
|
||||
# Check if this is a local file path or local file url
|
||||
exists = os.path.exists(path)
|
||||
|
||||
# Safe open validation
|
||||
if not self.safecheck(path):
|
||||
path = path if urlparse(path).scheme else os.path.realpath(path)
|
||||
raise IOError(f"Safeopen URL validation failed: path={path}, safeopen={self.safeopen}")
|
||||
|
||||
# Consider local files and HTTP urls valid
|
||||
return (path if exists or urlparse(path).scheme in ("http", "https") else None, exists)
|
||||
|
||||
def retrieve(self, url):
|
||||
"""
|
||||
Retrieves content from url.
|
||||
|
||||
Args:
|
||||
url: input url
|
||||
|
||||
Returns:
|
||||
data
|
||||
"""
|
||||
|
||||
# Local file
|
||||
if os.path.exists(url):
|
||||
with open(url, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
# Remote file
|
||||
return self.urlretrieve(url)
|
||||
|
||||
def safecheck(self, url):
|
||||
"""
|
||||
Safe open url validation. Validates a local path is within the safeopen directory and
|
||||
that URLs are a public HTTP(s) URLs.
|
||||
|
||||
Args:
|
||||
url: input url
|
||||
|
||||
Returns:
|
||||
True if url is valid, false otherwise
|
||||
"""
|
||||
|
||||
# Default to allow all urls when safe open is disabled
|
||||
valid = True
|
||||
|
||||
if self.safeopen:
|
||||
if os.path.exists(url) or not self.allowurl:
|
||||
# Validate local file is in safe path
|
||||
path = os.path.realpath(url)
|
||||
prefix = os.path.commonpath([self.safeopen, path])
|
||||
valid = prefix == self.safeopen
|
||||
else:
|
||||
# URL validation
|
||||
valid = url.lower().startswith("http") and not self.urlretrieve.isprivateurl(url)
|
||||
|
||||
return valid
|
||||
@@ -141,11 +141,11 @@ class Segmentation(Pipeline):
|
||||
elif self.sentences:
|
||||
content = [self.clean(x) for x in sent_tokenize(text)]
|
||||
elif self.lines:
|
||||
content = [self.clean(x) for x in re.split(r"\n{1,}", text)]
|
||||
content = [self.clean(x) for x in re.split(r"(?:\r?\n)+", text)]
|
||||
elif self.paragraphs:
|
||||
content = [self.clean(x) for x in re.split(r"\n{2,}", text)]
|
||||
content = [self.clean(x) for x in re.split(r"(?:\r?\n){2,}", text)]
|
||||
elif self.sections:
|
||||
split = r"\f" if "\f" in text else r"\n{3,}"
|
||||
split = r"\f" if "\f" in text else r"(?:\r?\n){3,}"
|
||||
content = [self.clean(x) for x in re.split(split, text)]
|
||||
else:
|
||||
content = self.clean(text)
|
||||
|
||||
@@ -99,12 +99,15 @@ class Tabular(Pipeline):
|
||||
columns = self.textcolumns
|
||||
if not columns:
|
||||
columns = list(df.columns)
|
||||
if self.idcolumn:
|
||||
if self.idcolumn is not None:
|
||||
columns.remove(self.idcolumn)
|
||||
|
||||
# Transform into (id, text, tag) tuples
|
||||
for index, row in df.iterrows():
|
||||
uid = row[self.idcolumn] if self.idcolumn else index
|
||||
# Tuple iteration preserves column types instead of coercing each row to one dtype.
|
||||
for index, values in zip(df.index, df.itertuples(index=False, name=None)):
|
||||
row = dict(zip(df.columns, values))
|
||||
# Coerce NaN ids to None so the row-index fallback applies.
|
||||
uid = self.column(row[self.idcolumn]) if self.idcolumn is not None else index
|
||||
uid = uid if uid is not None else index
|
||||
text = self.concat(row, columns)
|
||||
|
||||
@@ -112,10 +115,10 @@ class Tabular(Pipeline):
|
||||
|
||||
# Also add row for content
|
||||
if isinstance(self.content, list):
|
||||
row = {column: self.column(value) for column, value in row.to_dict().items() if column in self.content}
|
||||
row = {column: self.column(value) for column, value in row.items() if column in self.content}
|
||||
rows.append((uid, row, None))
|
||||
elif self.content:
|
||||
row = {column: self.column(value) for column, value in row.to_dict().items()}
|
||||
row = {column: self.column(value) for column, value in row.items()}
|
||||
rows.append((uid, row, None))
|
||||
|
||||
return rows
|
||||
@@ -135,7 +138,7 @@ class Tabular(Pipeline):
|
||||
parts = []
|
||||
for column in columns:
|
||||
column = self.column(row[column])
|
||||
if column:
|
||||
if column is not None and column != "":
|
||||
parts.append(str(column))
|
||||
|
||||
return ". ".join(parts) if parts else None
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user