Compare commits

...
306 Commits
Author SHA1 Message Date
Hua.T e981d2d8fc Forward backend options during streaming speech generation (#1389) (#1390) 2026-10-11 17:23:05 -04:00
PerryLink 55fadb4c56 Reject a semicolon in a SQL query instead of silently cutting the statement (#1388) 2026-10-11 17:21:01 -04:00
Rızgar Ozan fe5be99bfa Handle ServiceTask responses without a Content-Type or body (#1387) 2026-10-11 17:17:06 -04:00
David Mezzetti 243aaa191d Merge pull request #1385 from kartsan03/sqlite-load-switch
Switch SQLite ANN connections when load() is given a new path
2026-10-11 17:13:58 -04:00
leeshengis 93a767f004 Preserve workflow and task configuration dictionaries (#1377) 2026-10-11 17:12:41 -04:00
davidmezzetti 62099b9a34 Cleanup open connection in test 2026-10-11 12:03:14 -04:00
David Mezzetti c786f1b5ec Replace continue statement in HTMLToMarkdown list handler (#1357) (#1386) 2026-10-11 11:02:08 -04:00
amengnew 738df97d5f Fall back to row index when Tabular IDs are NaN (#1366) 2026-10-11 10:37:51 -04:00
Raphael Malikian 4605ee001f Fix LogOdds.calibrate to compute the mean once instead of per score (#1365) 2026-10-10 08:06:37 -04:00
Rızgar Ozan 46516fa4e6 Replace continue statement in HTMLToMarkdown list handler (#1357) 2026-10-10 14:24:22 +03:00
kartsan03 fe1fbbc933 Switch SQLite ANN connections when load() is given a new path 2026-10-10 11:11:24 +00:00
Hua.T 084a1eca4d Fix audio trimming of partial detection windows (#1374) (#1375) 2026-10-10 06:35:44 -04:00
Muneeb Ullah ba9d52310a Fix infertopics ignoring topiced neighbors when untopiced ones outnumber them (#1373) 2026-10-10 06:33:12 -04:00
Jasper ca9563daac Use only the query argument of similar() clauses in terms and explain (#1367) 2026-10-09 10:29:46 -04:00
D41910 0e7202caee Honor an empty allow list in the bash agent tool (#1362) 2026-10-08 15:55:25 -04:00
David Mezzetti c1f1686889 Merge pull request #1361 from kartsan03/rdbms-graph-addattribute
Write RDBMS graph attribute changes to the database
2026-10-08 13:06:57 -04:00
kartsan03 04f5bf1998 Write RDBMS graph attribute changes to the database 2026-10-08 12:52:02 +00:00
Gaurav Chintakunta c8e7d813af Preserve ordered-list start values in HTMLToMarkdown (#1355) (#1356) 2026-10-08 07:05:07 -04:00
Yuilona edf61d1d4a Use each query's hybrid weights when batching similar clauses (#1359) 2026-10-07 16:21:36 -04:00
Raphael Malikian 7e92f5a7d7 Fix StorageTask listing failing on tagged storage inputs (#1348) 2026-10-07 14:15:17 -04:00
Raashish Aggarwal b2ff9bd85c Sort NULL values like SQLite when merging ordered results (#1353) 2026-10-07 13:09:08 -04:00
Raashish Aggarwal 9c718f7876 Merge aggregate columns written in uppercase across shards (#1352) 2026-10-07 09:00:39 -04:00
Moinuddin Shaik c486f4706e Return no results from Reranker when a search has no matches (#1350) 2026-10-07 09:00:03 -04:00
Lukas 341184ba46 Keep word spacing and drop empty link targets in HTMLToMarkdown (#1347) 2026-10-07 08:58:51 -04:00
Lukas 647d80f95e Replace HTML line breaks with newlines in HTMLToMarkdown (#1345) 2026-10-07 07:22:03 -04:00
Yash Raj Pandey 408c25d201 Explain DuckDB query errors that abort the transaction (#1341) 2026-10-07 06:34:30 -04:00
davidmezzetti b0f29a11d1 Bump Docker version 2026-10-06 10:18:15 -04:00
davidmezzetti 49a64943b8 Bump version 2026-10-06 10:18:02 -04:00
David Mezzetti af33f2e024 Merge pull request #1333 from HuaTNA/fix-tabular-falsy-idcolumn
Respect explicitly selected falsy ID column names in Tabular
2026-10-05 11:06:24 -04:00
davidmezzetti f908b5d88d Support trust_remote_code in associated AutoModels, closes #1342 2026-10-05 10:25:41 -04:00
David Mezzetti 2e8493dfca Merge pull request #1338 from serhiizghama/fix/duckdb-bind-parameter-order
Keep DuckDB bind values in query order when a parameter repeats
2026-10-05 08:53:15 -04:00
shabeeth2 d7e3e29ac3 docs: match task input tuple in explanation (#1337) 2026-10-05 06:22:06 -04:00
Mohammad Hijjawi 4e2efa2cbc Keep node id 0 in graph path search results (#1335) 2026-10-05 06:19:12 -04:00
Gaurav Chintakunta ffcbc0b815 Skip NULL aggregate values from shards with no matching rows (#1331) 2026-10-05 06:17:05 -04:00
tuanzirwar f14c26c4c7 Fix ordinary streaming output buffering with stripthink enabled (#1311) 2026-10-05 06:15:44 -04:00
serhiizghama 346bb9d282 Test a repeated bind parameter followed by another parameter 2026-10-05 11:12:23 +07:00
serhiizghama e3e0413747 Keep DuckDB bind values in query order when a parameter repeats 2026-10-05 10:47:53 +07:00
David Mezzetti 32655393e0 Merge pull request #1323 from Lesereingrape/fix/concat-merge-falsy-values-1003
Keep 0 and False in concat merge outputs
2026-10-04 11:08:45 -04:00
davidmezzetti 9118875980 Move coverage check to only run for Linux GitHub action 2026-10-04 09:45:08 -04:00
Manobhi Sriram d9b4476187 Decode bash tool output as UTF-8 instead of locale encoding, closes #1326 (#1327) 2026-10-04 09:32:30 -04:00
Grapette.L ff5aec468a Fix DuckDB bind parameters followed by a non-word character not being converted (#1325) 2026-10-04 09:01:16 -04:00
davidmezzetti 3446a4ef13 Fail if code coverage less than 100% 2026-10-04 08:57:07 -04:00
HuaTNA 84aa1c81c7 Fix Tabular handling of explicitly selected falsy ID columns (#1332) 2026-10-03 19:14:23 -04:00
Lesereingrape 6bc76132d0 Keep 0 and False in concat merge outputs 2026-10-03 11:34:59 +08:00
Morgan Carr 893d0424be Cache document token vectors in Reranker (#1321) 2026-10-02 14:32:17 -04:00
Edson 98c965cc3b Keep all rows for native vstack merges with list inputs (#1320) 2026-10-02 14:30:31 -04:00
Yash Raj Pandey 00622f2cf1 Parse true and false as boolean literals in SQL queries (#1316) 2026-10-02 14:28:51 -04:00
Yash Raj Pandey 55d52fda92 Keep graph categories aligned with topics after a delete (#1315) 2026-10-02 14:24:59 -04:00
Ni0ri 20c9b8e0be Add Zvec Sparse ANN (#1307) 2026-10-01 09:53:38 -04:00
Mohammad Hijjawi 3f33fb2918 Format nested HTML lists without duplicating items (#1286) 2026-10-01 06:48:14 -04:00
davidmezzetti 3c7a548923 Remove build script workaround, closes #1239 2026-09-30 13:03:03 -04:00
Pablo Miguel Argudo c78fbd976c Support append and delete in the RaBitQ ivf mode (#1304) 2026-09-30 12:44:19 -04:00
David Mezzetti 2d535ff215 Merge pull request #1303 from kartsan03/fix/ann-delete-numpy-ids
Delete the requested rows when ANN ids are a NumPy array
2026-09-30 12:40:40 -04:00
kartsan03 6765330f4a Delete the requested rows when ANN ids are a NumPy array 2026-09-30 15:26:24 +00:00
davidmezzetti 1d5041b6e7 Fix test merge error 2026-09-30 11:07:02 -04:00
davidmezzetti d467878064 Refactor ANN tests, closes #1305 2026-09-30 11:04:07 -04:00
davidmezzetti 3bc07b3f09 Remove unnecessary note 2026-09-30 11:03:37 -04:00
Hua.T 8ef454d0bf Preserve keyword document ID types across save/load (#1274) (#1275) 2026-09-30 10:57:41 -04:00
Jin 47d3e65756 Fix training with pandas DataFrames, closes #1265 (#1266) 2026-09-30 10:55:44 -04:00
LimbC-C 81f908874b Open lazy SQLite connections before saving indexes (#1290) 2026-09-30 09:30:14 -04:00
David Mezzetti bed91a0b9f Merge pull request #1301 from Bdysj/openai-compliant-chat-completions
Return OpenAI-compliant index, created and finish_reason from chat completions
2026-09-30 09:05:50 -04:00
David Mezzetti 7df50abb61 Merge pull request #1282 from LimbC-C/fix/sqlite-binary-scores
Normalize SQLite binary ANN similarity scores
2026-09-30 09:03:45 -04:00
David Mezzetti 80b8ceff4a Merge pull request #1279 from LimbC-C/fix/array-negative-delete-ids
Ignore negative deletion IDs in NumPy and Torch ANN indexes
2026-09-30 08:59:03 -04:00
David Mezzetti de9e09e9ae Merge pull request #1299 from LimbC-C/fix/tabular-numeric-types
Preserve integer document IDs and numeric column types in Tabular
2026-09-30 08:19:59 -04:00
Morgan Carr 69ede9ec3e Fall back to root tokenizers for ONNX and LiteRT Hub subdirectory models (#1277) 2026-09-30 07:41:03 -04:00
Pablo Miguel Argudo f8689ed903 Store only the quantized index in the RaBitQ backend (#1297) 2026-09-30 07:05:37 -04:00
bob 0abf3763b4 Return OpenAI-compliant index, created and finish_reason from chat completions, closes #1300 2026-09-30 14:51:02 +08:00
LimbC-C 11b2cfe298 Preserve numeric column types in Tabular rows (#1298) 2026-09-30 09:45:58 +08:00
LimbC-C d0cc2bed9c Simplify regression test in existing TestDense class (#1282) 2026-09-30 09:39:23 +08:00
LimbC-C a5ea35d7a7 Simplify regression test in existing TestDense class (#1279) 2026-09-30 09:39:18 +08:00
Yash Raj Pandey 6a15ed8490 Include offset rows in similar() candidates (#1291) 2026-09-29 17:25:52 -04:00
David Mezzetti 76147369fb Merge pull request #1295 from Bdysj/pad-summary-batches
Pad batched inputs in the Summary pipeline
2026-09-29 15:27:58 -04:00
bob 2e89af0845 Pad batched inputs in the Summary pipeline, closes #1294 2026-09-30 01:19:38 +08:00
davidmezzetti d9d6a1cbb1 Update pooling mode check 2026-09-29 11:05:04 -04:00
davidmezzetti f87574de4c Support pooling_mode with models trained with recent Sentence Transformers, closes #1292 2026-09-29 10:23:16 -04:00
Mohammad Hijjawi e313248ff2 Keep uppercase tokens in alphanumeric tokenization with lowercase=False (#1287) 2026-09-29 09:03:03 -04:00
davidmezzetti 5cba72c88a Update pylintrc file and remove warning overrides for new limits 2026-09-29 07:11:03 -04:00
davidmezzetti 45d334f4ce Remove unnecessary new test file and consolidate test into existing class 2026-09-29 07:10:09 -04:00
Jin 95d4c1c881 Label QA spans without the answer with the CLS token, closes #1263 (#1264) 2026-09-29 06:36:47 -04:00
LimbC-C ef28da8735 Reject iterator inputs for multi-action tasks without buffering (#1272) 2026-09-28 10:40:36 -04:00
LimbC-C f1304d8403 Normalize SQLite binary ANN similarity scores (#1281) 2026-09-28 09:43:12 +08:00
davidmezzetti 97ee7e073a Bump job timeout given that the Windows build frequently times out 2026-09-27 05:51:01 -04:00
LimbC-C 1d3f8e9ebf Ignore negative NumPy and Torch ANN deletion IDs (#1278) 2026-09-26 09:43:32 +08:00
davidmezzetti f3ceece76a Add test for missing SQLite ANN coverage 2026-09-25 09:18:02 -04:00
Pablo Miguel Argudo 6737ee4e01 Add RaBitQ as an ANN backend (#1251) 2026-09-25 08:39:39 -04:00
davidmezzetti e30e8edbcc Add test for missing SQLite ANN coverage 2026-09-25 08:00:21 -04:00
davidmezzetti 483dbd3b3b Update test 2026-09-24 21:49:06 -04:00
davidmezzetti c09a6f24f1 Update test 2026-09-24 21:46:09 -04:00
davidmezzetti 6e7f0d260e Update test 2026-09-24 21:12:15 -04:00
davidmezzetti 793ecc8dc2 Require Python 3.11, closes #1273 2026-09-24 16:58:32 -04:00
David Mezzetti 37500887ae Merge pull request #1268 from HuaTNA/fix-rag-input-mutation
Preserve reusable RagTask inputs
2026-09-24 16:35:58 -04:00
David Mezzetti 4e2eb242b6 Support explicit LIKE escapes in wildcard term queries (#1259) 2026-09-24 16:34:16 -04:00
Jin 83c75ea28a Pack rows into chunks up to and including maxlength, closes #1269 (#1270) 2026-09-24 16:31:37 -04:00
LimbC-C f47d389735 Preserve wildcard compatibility and support explicit LIKE escapes 2026-09-24 23:54:33 +08:00
HuaTNA 4bb2141954 Preserve reusable RagTask inputs (#1267) 2026-09-23 23:48:46 -04:00
David Mezzetti 5d8c999ae0 Merge pull request #1262 from HuaTNA/fix-template-falsy-rules
Preserve falsy TemplateTask rule matches
2026-09-23 16:36:47 -04:00
HuaTNA 48a34e21c3 Preserve falsy TemplateTask rule matches (#1261) 2026-09-23 13:22:36 -04:00
davidmezzetti 262b03027e Code cleanup 2026-09-23 08:01:47 -04:00
刘冠杰 36c375b86e Honor the nfeatures setting in IVFSparse (#1250) 2026-09-23 07:53:18 -04:00
David Mezzetti b715812e68 Merge pull request #1248 from LimbC-C/fix/deleted-keyword-results
Keep deleted keyword results excluded after common-term scoring
2026-09-23 07:51:27 -04:00
hylin dcdded892d Keep community nodes that match none of the topic terms (#1243) 2026-09-23 07:40:25 -04:00
LimbC-C e7894cdb00 Escape literal SQL wildcard characters in term expansion (#1258) 2026-09-23 13:06:27 +08:00
LimbC-C 9f9b32e070 Clear deleted scores at the end of common-term merging 2026-09-23 13:00:30 +08:00
David Mezzetti 00167b9f08 Merge pull request #1256 from dingpuyu/codex/fix-retrieve-local-filenames
Preserve local filenames in RetrieveTask
2026-09-22 14:00:19 -04:00
dingpuyu 3c07766199 Preserve local filenames in RetrieveTask 2026-09-23 00:36:08 +08:00
Amir Fathi 7cb6bd09db Fix DuckDB.formatargs dropping repeated named bind parameters (#1253) 2026-09-22 12:22:45 -04:00
Caner Olcay 4b7d70e9e8 Keep SQLite ANN data when saving a modified index to a new path (#1249) 2026-09-22 06:55:29 -04:00
David Mezzetti 2ed4e46ebf Merge pull request #1245 from serhiizghama/fix/cluster-forward-weights-and-encode-parameters
fix(api): forward zero weights and url-encode parameters in Cluster search
2026-09-22 06:54:14 -04:00
LimbC-C 47d5f56768 Keep deleted keyword results excluded after common-term scoring (#1247) 2026-09-19 09:46:09 +08:00
serhiizghama d5db643540 test: cover cluster search argument forwarding 2026-09-18 09:15:19 +07:00
serhiizghama 9f7b4ddd42 fix(api): forward zero weights and url-encode parameters in Cluster search 2026-09-18 08:56:12 +07:00
Morgan Carr 19ff237c37 Set the HNSW delete counter when loading an index (#1242) 2026-09-15 14:21:46 -04:00
Morgan Carr ad77514b32 Store IVFSparse deletes as int (#1241) 2026-09-15 14:20:36 -04:00
davidmezzetti 144a957d42 Test latest version of LiteLLM #1193 2026-09-15 07:51:21 -04:00
Serhii Zghama 2e8a3d7892 fix(database): don't drop rows when sorting or grouping by a JSON field (#1238) 2026-09-14 10:19:52 -04:00
Amir Fathi 72f9ca7605 Fix Cluster combining avg() as an unweighted mean of shard averages (#1237) 2026-09-14 10:18:52 -04:00
Amir Fathi ebe13ae304 Fix quantize: false silently enabling Q8_0 quantization in the GGML backend (#1236) 2026-09-11 08:25:40 -04:00
Serhii Zghama 5d0e995f7f fix(ann): honor quantize: false in the Faiss and SQLite backends (#1234) 2026-09-11 08:22:49 -04:00
davidmezzetti 7965a7f42e Build script workaround #1239 2026-09-10 08:52:56 -04:00
Amir Fathi 2bae6e4d43 Fix Transform.stream() yielding a None text value into the encoder (#1231) 2026-09-08 09:44:53 -04:00
Morgan Carr 767272f2d5 Key the vectors model cache by configuration (#1223) 2026-09-04 16:47:49 -04:00
Morgan Carr 21460347dd Fix ANN delete on quantized Torch stores and tiny GGML indexes (#1228) 2026-09-04 09:54:42 -04:00
Morgan Carr f3945f9120 Select the top documents in LateEncoder by the document axis (#1227) 2026-09-04 09:52:40 -04:00
Paco Cartones 3a7b597d78 fix(workflow): add request timeout to ServiceTask (#1226) 2026-09-04 09:00:49 -04:00
Morgan Carr 20e2dc41e1 Save empty keyword and content indexes without errors (#1224) 2026-09-04 08:21:11 -04:00
Morgan Carr eafe574d92 Keep the archive working directory alive when saving a loaded index (#1222) 2026-09-04 08:19:36 -04:00
Morgan Carr 7e98746d78 Delete every index entry for an id in Terms (#1219) 2026-09-04 08:11:47 -04:00
Morgan Carr 6ebfbbb716 Improve late pooling configuration errors (#1217) 2026-09-03 13:52:46 -04:00
David Mezzetti 8fef6ac454 Merge pull request #1220 from serhiizghama/fix/sparse-vectors-remove-spool-file
fix(vectors): remove the sparse spool file when checkpointing is disabled
2026-09-03 13:05:09 -04:00
davidmezzetti 0eafc0580f Update vectors package to standardize fallback logic, closes #1221 2026-09-03 11:53:55 -04:00
serhiizghama 5ade7881c8 test: cover spool file cleanup for sparse vectors 2026-09-03 20:19:31 +07:00
serhiizghama 948dc42d19 fix(vectors): remove the sparse spool file when checkpointing is disabled 2026-09-03 19:58:47 +07:00
David Mezzetti 2e8af221f4 Keep count() accurate after repeated deletes in GGML and IVFSparse (#1218)
Deleting an id that is already marked deleted no longer decrements count()
again in the GGML and IVFSparse ANN backends. IVFSparse also ignores ids that
were never indexed, matching the dense backends.
2026-09-02 18:18:11 -04:00
Morgan Carr 396b3b8065 Resolve label filters by label id (#1216) 2026-09-02 17:15:45 -04:00
Morgan Carr 80420ee10a Ignore padding rows in LateEncoder scores (#1215) 2026-09-02 17:14:39 -04:00
David Mezzetti 1d45e88569 Merge pull request #1214 from morgan-coded/fix/milvus-load-collection
Load the Milvus collection when opening a saved index
2026-09-02 15:25:19 -04:00
Joe Bergin 007d61077a Add ONNX vectors backend (#1212) 2026-09-02 12:19:57 -04:00
David Mezzetti 84fcd81039 fix(api): parse the graph query parameter as a bool (#1213) 2026-09-02 11:18:16 -04:00
morgan-coded cfe2227f30 Load the Milvus collection when opening a saved index
With milvus-lite 3.2.1, a reopened collection that has an index starts in the
released state, so search failed after save and after load while count worked.
2026-09-02 03:20:29 -05:00
morgan-coded 82b5320da3 Keep count() accurate after repeated deletes in GGML and IVFSparse
Deleting an id that is already marked deleted no longer decrements count()
again in the GGML and IVFSparse ANN backends. IVFSparse also ignores ids that
were never indexed, matching the dense backends.
2026-09-02 02:53:52 -05:00
serhiizghama 18b93be702 test: cover graph parameter parsing in the search API 2026-09-02 11:06:18 +07:00
serhiizghama 027c07deb9 fix(api): parse the graph query parameter as a bool 2026-09-02 10:53:23 +07:00
Paco Cartones 876cd07cf7 Fix bash tool hanging on commands that read stdin (#1210) 2026-09-01 20:04:29 -04:00
davidmezzetti 546e7777f1 Bump version 2026-08-27 12:48:36 -04:00
Yash Raj Pandey 27542ca64b Index zero and False values in the tabular pipeline (#1207) 2026-08-27 07:39:36 -04:00
David Mezzetti f0ee7d2b2d Merge pull request #1206 from morgan-coded/fix/nprobe-clamp
Clamp default Faiss nprobe to nlist
2026-08-26 08:23:31 -04:00
morgan-coded bba7fd589e Clamp default nprobe to nlist 2026-08-25 17:18:46 -05:00
David Mezzetti 6349ca2eac Merge pull request #1205 from Evhye38496/fix-safetensors-import
Fix safetensors numpy/torch submodule imports
2026-08-25 07:58:29 -04:00
Evhye38496 ba1693480e Fix safetensors numpy/torch submodule imports #1204 2026-08-25 12:57:51 +08:00
Sainikhil Juluri e90c67ef06 fix(scoring): keep count() accurate when an id is deleted more than once (#1202) 2026-08-24 21:55:11 -04:00
Amir Fathi 57991cb164 Fix WordVectors.index() silently dropping the checkpoint argument (#1198) 2026-08-24 20:38:28 -04:00
David Mezzetti c2b264d297 Merge pull request #1201 from serhiizghama/fix/numpy-hamming-search-skip-deleted-rows
fix(ann): exclude deleted rows from quantized numpy search results
2026-08-24 15:57:28 -04:00
serhiizghama 406e417be0 test(ann): cover deleted-row exclusion for quantized numpy backend 2026-08-24 17:18:26 +07:00
serhiizghama 954b17af4d fix(ann): exclude deleted rows from quantized numpy search results 2026-08-24 17:05:58 +07:00
davidmezzetti 2a139449dc Re-add LiteLLM workaround #1193 2026-08-23 19:08:33 -04:00
davidmezzetti a10667a1c2 Test new version of LiteLLM #1193 2026-08-22 22:10:01 -04:00
davidmezzetti 34e6efcdd6 Update docs 2026-08-22 22:09:34 -04:00
David Mezzetti 9b7a96b4d1 Resolve LEMUR and MUVERA settings (#1196) 2026-08-22 21:49:14 -04:00
morgan-coded b94f770301 Resolve LEMUR and MUVERA settings 2026-08-22 16:19:06 -05:00
Arnaud Thery adbcccda46 Merge pull request #1195 from St4r4x/fix/crossencoder-multiclass-labels
fix: rank CrossEncoder by a fixed label, not top-1 position
2026-08-22 14:15:21 -04:00
David Mezzetti 20f818f72c Merge pull request #1194 from morgan-coded/fix/1190-lemur-collection-centering
Fix LEMUR collection centering
2026-08-19 06:37:51 -04:00
morgan-coded ee0416e87b Add centermean validation test 2026-08-18 20:22:07 -05:00
morgan-coded 00ec81a10e Fix LEMUR collection centering 2026-08-18 10:54:10 -05:00
davidmezzetti 059ecec2ee Build script workaround #1193 2026-08-18 09:58:35 -04:00
davidmezzetti 21c2c38b98 Build script workaround #1193 2026-08-18 09:56:00 -04:00
davidmezzetti 47c05994bf Cleanup of Pooling PRs 2026-08-18 09:53:29 -04:00
Kole Harvey 8f3735101b Resolve pooling method from a config missing pooling_mode_mean_tokens 2026-08-18 09:14:20 -04:00
Aryan Pardeshi 438ade31b6 Add max pooling 2026-08-18 09:12:32 -04:00
davidmezzetti 3ea197468d Code cleanup 2026-08-12 09:35:13 -04:00
LHMQ878 99ebbdfea2 fix: create default agent tools on first use (#1175) 2026-08-12 07:30:37 -04:00
David Mezzetti 5570dbc7b6 Merge pull request #1183 from FU-max-boop/fix/openai-latest-message
Fix OpenAI chat message selection
2026-08-12 07:29:43 -04:00
FU-max-boop 01b456900a Fix OpenAI chat message selection 2026-08-11 13:26:29 +08:00
davidmezzetti a0857e980a Update output validation #1181 2026-08-10 22:02:42 -04:00
davidmezzetti 84be3476d8 Add safeopen validation to retrieve task output, closes #1181 2026-08-10 19:29:05 -04:00
davidmezzetti f35d932979 Refactor late encoding logic to cleanup logic paths, closes #1180 2026-08-10 18:47:58 -04:00
Tai An a1c9b94d00 fix(muvera): consume true token rows instead of batch-padded ones (#1177) 2026-08-10 13:24:21 -04:00
Aryan Pardeshi 11ebacdb23 Convert MUVERA algorithm from NumPy to Torch (#1178) 2026-08-10 13:20:58 -04:00
davidmezzetti 5986fa6082 Test Docling build #1176 2026-08-10 10:57:48 -04:00
davidmezzetti 893a333887 Test Docling build fix #1176 2026-08-10 10:15:16 -04:00
David Mezzetti e54255f4e4 Merge pull request #1179 from latent-9/fix-spelling-typos
Fix spelling typos in documentation and docstrings
2026-08-10 09:08:58 -04:00
latent-9 f89835bc77 Fix spelling typos in documentation and docstrings 2026-08-10 04:40:37 +12:00
davidmezzetti ea531b3420 Build script workaround #1176 2026-08-04 12:24:05 -04:00
davidmezzetti 756917f14b Build script workaround #1176 2026-08-04 12:20:17 -04:00
davidmezzetti b1f4845284 Build script workaround #1176 2026-08-04 12:13:14 -04:00
David Mezzetti d2b1d3634f Merge pull request #1172 from morgan-coded/feat/lemur-trainer
Consolidate LEMUR training in the trainer pipeline
2026-08-04 10:12:05 -04:00
morgan-coded 193facd144 test: cover the tqdm library stub 2026-08-03 20:50:04 -05:00
morgan-coded 436662e32b feat: promote tqdm to a default library dependency 2026-08-03 14:32:12 -05:00
morgan-coded 54644a0526 chore: add lint pragma and docstring for relocated trainer tests 2026-08-02 18:11:54 -05:00
morgan-coded 817a2b142a fix: keep minimal installs importable 2026-08-02 17:04:51 -05:00
morgan-coded 987cff001f feat: consolidate LEMUR training in trainer
Move fitting, data preparation, model setup, epoch selection and progress feedback into LemurTrainer while keeping artifact loading and encoding in the pooler.
2026-08-02 16:55:32 -05:00
David Mezzetti 8383e6ce8d Merge pull request #1171 from morgan-coded/chore/lemur-style
Align LEMUR names with txtai style
2026-08-02 16:49:32 -04:00
morgan-coded 0be1231d8d chore: align LEMUR names with txtai style 2026-08-02 14:21:38 -05:00
David Mezzetti 79e277fb7d Merge pull request #1168 from morgan-coded/feat/center
Add configurable mean centering to late pooling
2026-08-02 13:36:13 -04:00
morgan-coded f4472312ba test: use portable tolerances for cross-batch comparisons 2026-08-02 06:47:56 -05:00
morgan-coded 90f7cc2078 test: isolate uncentered LEMUR coverage 2026-08-01 22:40:57 -05:00
morgan-coded ce11b08ab1 fix: address centering review feedback 2026-08-01 22:13:38 -05:00
morgan-coded 75b4ce5c7f feat: add configurable late pooling centering 2026-08-01 18:37:42 -05:00
David Mezzetti 2ef401bb78 Merge pull request #1164 from morgan-coded/feat/lemur
Add LEMUR fixed dimensional encoding
2026-08-01 17:35:07 -04:00
davidmezzetti 8c43c70687 Bump version 2026-07-30 17:43:11 -04:00
David Mezzetti 48e22b593c Merge pull request #1163 from AmirF194/fix/scoring-delete-unindexed-id
Fix Scoring.delete() crashing on ids that were never indexed
2026-07-30 14:48:11 -04:00
davidmezzetti 3beea7f4bf Add safeopen option for retrieve task, closes #1166 2026-07-30 14:36:52 -04:00
morgan-coded 626cc8fbe0 Add LEMUR validation coverage 2026-07-30 12:03:07 -05:00
davidmezzetti 21b9049c86 Allow disabling API routes, closes #1165 2026-07-30 12:49:38 -04:00
davidmezzetti abf504f70e Update configuration section for workflow tasks 2026-07-30 12:44:34 -04:00
davidmezzetti d9e9a9e684 Pin mcp < 2.0 until upstream issues resolved #1161 2026-07-30 12:43:47 -04:00
Amir Fathi bd199158b6 Fix Scoring.delete() crashing on ids that were never indexed 2026-07-30 12:35:35 +00:00
morgan-coded 390cc08b11 Document LEMUR training and configuration 2026-07-30 07:02:16 -05:00
morgan-coded feac4c8306 Test LEMUR pooling and training 2026-07-30 07:02:16 -05:00
morgan-coded d7f974b3d7 Add LEMUR training pipeline 2026-07-30 07:02:16 -05:00
morgan-coded d8fa848f73 Integrate LEMUR late pooling 2026-07-30 07:02:16 -05:00
morgan-coded bbd23bb12c Add LEMUR fixed dimensional encoding 2026-07-30 07:02:16 -05:00
David Mezzetti 8362110c64 Merge pull request #1162 from morgan-coded/fix/milvus-lite-version-floor
Set minimum milvus-lite version
2026-07-29 07:51:07 -04:00
morgan-coded 83e1835b6d Set minimum milvus-lite version 2026-07-28 20:33:31 -05:00
davidmezzetti cf2368269f Build script workaround #1161 2026-07-28 11:14:49 -04:00
davidmezzetti a910ddbbea Add Milvus tests on Windows #1156 2026-07-28 10:37:38 -04:00
davidmezzetti 7b49504dca Update test 2026-07-23 11:54:09 -04:00
davidmezzetti 86af7ec344 Test updates 2026-07-23 10:10:28 -04:00
Amir Fathi 50ca7b2fde Fix HNSW search RuntimeError when limit exceeds indexed element count (#1159) 2026-07-23 08:42:33 -04:00
David Mezzetti 7c8d001492 Merge pull request #1151 from link89/fix-openapi-auth-ui
fix Swagger UI auth issue
2026-07-23 08:30:09 -04:00
weihong.xu dea0a65fd1 implement auth with HTTPBearer for correct openapi spec 2026-07-23 09:56:26 +08:00
davidmezzetti a1224c06f1 Doc and test update 2026-07-22 09:43:36 -04:00
davidmezzetti 162b65d3d9 Remove file 2026-07-22 07:38:17 -04:00
davidmezzetti 363f8ca7cf Ensure explain / batchexplain don't return numpy values, closes #1157 2026-07-22 06:53:07 -04:00
davidmezzetti 1d084941f4 Code cleanup 2026-07-22 06:37:42 -04:00
Andrew Chen 597ba8133d Fix msgpack Accept header crashing streaming API endpoints (#1144) 2026-07-22 06:14:07 -04:00
Andrew Chen 6bbb25872a Fix Terms KeyError on a query with no similar() clause (#1142) 2026-07-22 06:13:33 -04:00
davidmezzetti f22bf5bbb5 Update test 2026-07-21 14:29:12 -04:00
davidmezzetti 83871c6fe3 Add documentation for Milvus ANN 2026-07-21 13:24:23 -04:00
davidmezzetti 2c6ec162d0 Skip Milvus tests on Windows #1156 2026-07-21 13:10:56 -04:00
davidmezzetti d59c88424d Code cleanup 2026-07-21 12:59:57 -04:00
David Mezzetti 95c7b8df1e Merge pull request #1149 from morgan-coded/milvus-lite-backend
Add milvus-lite embedded dense ANN backend
2026-07-21 12:41:05 -04:00
Loi Nguyen 2389b5946c Fix Segmentation dropping paragraph/section splits on CRLF line endings (#1152) 2026-07-21 12:29:12 -04:00
Andrew Chen b246b642c7 Fix LIMIT bind-parameter crashing content SQL search (#1148) 2026-07-21 12:26:15 -04:00
Andrew Chen 46467a3453 Raise SQLError on empty select expressions and a dangling AS (#1147) 2026-07-21 12:24:12 -04:00
davidmezzetti fcfd340ff1 Cleanup tests 2026-07-21 12:18:35 -04:00
Andrew Chen 77965f44dd Fix Cluster.shard TypeError on tuple documents and empty-string ids (#1146) 2026-07-21 11:57:57 -04:00
Andrew Chen 7435297bac Fix console printing a traceback for a query with no results (#1145) 2026-07-21 11:51:11 -04:00
Andrew Chen 3bb34eae2e Fix Terms.search AttributeError on an empty terms index (#1143) 2026-07-21 11:40:43 -04:00
Andrew Chen 79a28641b2 Fix TemplateTask.match KeyError when a rule key is missing from element (#1141) 2026-07-21 11:25:58 -04:00
Andrew Chen 7c77c134e2 Fix HTMLToMarkdown crash on a meta description tag with no content attribute (#1140) 2026-07-21 11:23:29 -04:00
Andrew Chen d6f696d795 Fix Aggregate IndexError on a SQL query that matched nothing across shards (#1139) 2026-07-21 11:19:28 -04:00
davidmezzetti d8ae6776f8 Only resolve transform functions when env variable set, closes #1155 2026-07-21 11:07:58 -04:00
davidmezzetti e18e3d1245 Reduce verbose comment 2026-07-21 10:40:32 -04:00
Andrew Chen ee0d1a0214 fix: Task.filteredrun crashes/corrupts output for single-action tasks with merge disabled (#1138) 2026-07-21 10:16:26 -04:00
davidmezzetti 5ddedf07cd Fix test order 2026-07-21 10:13:44 -04:00
David Mezzetti 590c71f593 Merge pull request #1154 from winklemad/fix/ivfsparse-sort-order
Fix IVFSparse returning unsorted sparse ANN search results
2026-07-21 10:11:17 -04:00
winklemad c33188a0f8 Fix IVFSparse returning unsorted sparse ANN search results
IVFSparse.topn selects the top n with np.argpartition but never sorts them, so
search results came back in arbitrary order instead of ranked by score descending
like every other backend (numpy argsort, hnsw/faiss, sqlite ORDER BY, pgvector).
Sort each row by score after the partition. The LateEncoder pipeline used the same
argpartition-without-sort pattern in its ranking path and is fixed the same way.
2026-07-21 18:46:36 +05:30
davidmezzetti 638a2e12b1 Validate resolver target classes, closes #1153 2026-07-21 09:13:43 -04:00
morgan-coded 03da5c96fb Rename milvuslite backend to milvus per maintainer request 2026-07-17 08:08:57 -05:00
morgan-coded 99086302ce Add milvus-lite as a dense ANN backend 2026-07-16 11:15:18 -05:00
davidmezzetti 2f11f93743 Code cleanup 2026-07-14 15:27:55 -04:00
davidmezzetti 783d9cf680 Cleanup docs 2026-07-14 15:15:21 -04:00
David Mezzetti 281babeff7 Merge pull request #1137 from chuenchen309/fix/graph-filter-unknown-node-id
Fix Graph.filter() crash on ids that don't exist in the graph
2026-07-14 14:19:06 -04:00
David Mezzetti a71732198a Merge branch 'master' into fix/graph-filter-unknown-node-id 2026-07-14 14:18:14 -04:00
Morgan Carr 14cbe51e04 Add zvec as a dense ANN backend (#1130) 2026-07-14 13:41:35 -04:00
davidmezzetti 5d15f7fa66 Revert #1136 as it's confirmed to not be necessary 2026-07-14 13:40:09 -04:00
davidmezzetti ac2dc6a30b Code cleanup 2026-07-14 13:36:50 -04:00
chuenchen309andClaude Opus 4.8 7791174c1b Fix Graph.filter() crash on ids that don't exist in the graph
Graph.filter() unpacked self.node(node) with ** into addnode() for every
input id, but self.node() returns None for an id not present in the
graph. **None raises TypeError instead of the filter simply skipping an
id it doesn't recognize - the docstring says nodes are "nodes to select",
with no stated precondition that every id must already exist.

Skip ids where self.node() returns None instead of crashing.

Co-Authored-By: Claude Opus 4.8 <[email protected]>
2026-07-15 00:09:32 +08:00
David Mezzetti b21ba7065e Merge pull request #1136 from chuenchen309/fix/graph-insert-none-text-index-desync
Fix Graph.insert() node id desync when a document's text/object field is None
2026-07-14 11:11:06 -04:00
chuenchen309andClaude Opus 4.8 10a6b23304 Fix Graph.insert() node id desync when a document's text/object field is None
Graph.insert() decided whether a document consumed an indexid slot based
on whether the resolved text/object value was not None. But the actual
vector index offset, computed by Transform.stream() (embeddings/index/
transform.py), increments based on key presence (`self.text in
document[1]`), regardless of the value. A document like {"text": None}
consumes a real vector index slot but was skipped by Graph.insert()
without incrementing its own index counter, so every subsequent document
in the batch got assigned to the wrong graph node id - a silent,
permanent misattribution of node ids to the wrong document data.

Track slot consumption with the same key-presence criterion Transform.
stream() uses (hasfield), separately from whether a node actually gets
created for that slot (still gated on the resolved value being non-None,
so no node is created for a genuinely empty document - but the index
still advances, leaving a gap instead of shifting every later id).

Co-Authored-By: Claude Opus 4.8 <[email protected]>
2026-07-14 22:55:12 +08:00
David Mezzetti 576f0ddcd7 Merge pull request #1135 from chuenchen309/fix/filetask-urltask-non-string-crash
Fix FileTask/UrlTask crash on non-string elements
2026-07-14 10:36:26 -04:00
chuenchen309andClaude Opus 4.8 9b8a606a16 Fix FileTask/UrlTask crash on non-string elements
Task.accept() is deliberately written to tolerate arbitrary element types
(including None) so heterogeneous batches can flow through a multi-task
workflow, letting a task pass through elements it doesn't recognize rather
than crashing. UrlTask.accept and FileTask.accept override this but called
element.lower()/re.sub(..., element) unconditionally, assuming element was
already a string:

- UrlTask.accept: super().accept(element) short-circuits True whenever no
  select filter is configured, then element.lower() raises AttributeError
  on any non-string element (including None).
- FileTask.accept: re.sub(FileTask.FILE, "", element) ran before the
  isinstance(element, str) check further down, so that check was dead code
  for the exact case it looked like it guarded against, and raised
  TypeError on any non-string element.

Both now check isinstance(element, str) before doing string operations, so
non-string/None elements pass through unchanged instead of crashing.
ImageTask inherits FileTask.accept and needed no change - it's protected
by the same fix via short-circuit evaluation.

Co-Authored-By: Claude Opus 4.8 <[email protected]>
2026-07-14 21:54:43 +08:00
davidmezzetti 01344b021c Test Python with build script #1134 2026-07-14 09:33:15 -04:00
davidmezzetti dd00608b9a Test Python with build script #1134 2026-07-14 09:28:05 -04:00
davidmezzetti 702e0b880f Test Python with build script #1134 2026-07-14 09:25:48 -04:00
davidmezzetti 14ac0e4cc5 Test Python 3.11 with build script #1134 2026-07-14 09:08:13 -04:00
davidmezzetti 1f28e5401d Test Python 3.11 with build script #1134 2026-07-14 08:48:00 -04:00
davidmezzetti a84f0a388c Fix test 2026-07-13 19:47:15 -04:00
davidmezzetti 510f9c18ed Update build script 2026-07-13 17:55:21 -04:00
davidmezzetti ae67f63359 Update build script 2026-07-13 17:39:07 -04:00
davidmezzetti 0e16a22051 Update build script 2026-07-13 17:10:17 -04:00
davidmezzetti bfd0aa7ff4 Update build script 2026-07-13 17:06:29 -04:00
davidmezzetti ba2ea2db4e Update build script 2026-07-13 16:57:19 -04:00
davidmezzetti 1daac761aa Update build script 2026-07-13 16:55:07 -04:00
davidmezzetti e9e65393a0 Update build script 2026-07-13 16:52:26 -04:00
davidmezzetti 24ca4e36ca Update build script 2026-07-13 16:32:59 -04:00
davidmezzetti 8208780b6a Update build script 2026-07-13 16:30:29 -04:00
davidmezzetti c3735dd5d4 Add tokenbatch size to HFTrainer pipeline, closes #1133 2026-07-13 16:15:02 -04:00
davidmezzetti 6544d123f6 Update test to work with latest version of pgvector 2026-07-13 16:08:53 -04:00
davidmezzetti aa0a8cfc32 Cleanup code from #1125 2026-07-02 10:37:57 -04:00
David Mezzetti fc6f93828b Merge pull request #1125 from Sanjays2402/fix/unterminated-sql-clause
Raise SQLError on unterminated bracket, function and similar clauses
2026-07-02 09:41:17 -04:00
Sanjay Santhanam 796c5bb0a8 Raise SQLError on unterminated bracket, function and similar clauses
The token consumer loops in Expression.bracket(), Expression.function() and
Expression.similar() read the next token with next(iterator, (None, None)) but
never checked for the end-of-stream sentinel. When a clause was left
unterminated (e.g. "select [a from txtai", "select func(a from txtai" or
"select * from txtai where similar('abc'"), the iterator was exhausted and the
code indexed tokens[None], raising a cryptic

    TypeError: list indices must be integers or slices, not NoneType

instead of the SQLError that the database layer documents and callers catch for
invalid SQL. The parse step runs before the query is handed to the database, so
this TypeError escaped the SQLError wrapping in Database.execute().

Add Expression.advance(), a small helper that wraps next() and raises SQLError
when the token stream is exhausted mid-clause, and route the three delimited
consumer loops through it. Auto-closing the clause instead was rejected because
it would silently reinterpret malformed input (e.g. "select [a" would become the
valid column "a"), masking the user error. This mirrors the existing None-guard
convention already used in Expression.compound().

Behavior on valid queries is unchanged (advance() only diverges from the old
code path when the iterator is empty). Adds testUnterminated covering
unterminated bracket, function and similar clauses.
2026-07-01 23:11:21 -07:00
davidmezzetti 7613acd6e5 Bump version 2026-07-01 15:24:55 -04:00
davidmezzetti 9db11953c5 Bump Torch version 2026-07-01 15:24:49 -04:00
davidmezzetti 2ecf4412fc Update build script 2026-07-01 10:32:49 -04:00
davidmezzetti 13b1a0f986 Update commonprefix to commonpath, closes #1120 2026-06-22 12:58:57 -04:00
davidmezzetti 89d5207b6e Limit tabular pipeline to local CSV files, closes #1119 2026-06-22 12:55:31 -04:00
davidmezzetti 21e6ab8e6a Add missing dependency #1118 2026-06-19 11:22:36 -04:00
davidmezzetti ae898e6905 Add text extraction with LiteParse, closes #1118 2026-06-19 10:51:58 -04:00
davidmezzetti 7aee5f91af Support multiple Dense layers for Late Interaction Pooling, closes #1117 2026-06-18 13:28:01 -04:00
davidmezzetti 7844e2c597 Fix how custom encoding is added to routers, closes #1115 2026-06-16 10:40:31 -04:00
davidmezzetti f9bfc2cb0d Update test 2026-06-16 10:39:37 -04:00
David Mezzetti 535edaf3ad Merge pull request #1114 from JSap0914/fix/ngram-tokenizer-lowercase
fix(tokenizer): respect lowercase=False in ngramtokenize
2026-06-16 07:20:07 -04:00
JSap0914 f507c442fd fix(tokenizer): respect lowercase=False in ngramtokenize
ngramtokenize called text.lower() unconditionally, which silently
overrode the lowercase=False setting. __call__ already lowercases text
before delegating to ngramtokenize when self.lowercase is True, so the
extra .lower() inside ngramtokenize is both redundant for the default
case and incorrect when lowercase=False.

Remove the redundant .lower() call and add a regression test.
2026-06-16 18:51:28 +09:00
davidmezzetti 73e2f6ee51 Remove find from the default list of bash tools, closes #1113 2026-06-11 12:09:45 -04:00
davidmezzetti bc8847edb0 Set minimum transformers version 2026-06-11 11:30:50 -04:00
davidmezzetti d69b72492d Test Transformers required version 2026-06-11 09:59:24 -04:00
davidmezzetti e0ae1450b4 Use gliner fork to relax transformers version caps, closes #1112 2026-06-11 09:08:05 -04:00
davidmezzetti 24c20cc0ff Revert noisy logging workaround, closes #1106 2026-06-11 09:06:59 -04:00
davidmezzetti 5f233d4a74 Update test 2026-06-09 17:16:41 -04:00
davidmezzetti 11b32da720 Add reindex parameter for API, closes #1111 2026-06-09 15:32:04 -04:00
davidmezzetti 039568541b Add note on ALLOW_PICKLE 2026-06-09 15:31:12 -04:00
davidmezzetti df53803cf7 Add turbovec an ANN backend, closes #1109 2026-06-08 12:18:28 -04:00
davidmezzetti 3c4d32f8f7 Migrate from ggml-py to ggml-python #1105 2026-06-08 07:02:18 -04:00
davidmezzetti 68e1c5719e Bump version 2026-06-04 10:57:45 -04:00
davidmezzetti c3211fbce9 Bump Torch version 2026-06-04 10:57:39 -04:00
221 changed files with 9458 additions and 1407 deletions
+8 -4
View File
@@ -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,7 +40,7 @@ jobs:
echo "TIKA_STARTUP_SLEEP=30" >> $GITHUB_ENV
echo "TIKA_STARTUP_MAX_RETRY=10" >> $GITHUB_ENV
brew install portaudio
sudo xcode-select -s "/Applications/Xcode_16.app"
pip install torch==2.12.1 torchvision==0.27.1
if: matrix.os == 'macos-latest'
- name: Install dependencies - Windows
@@ -61,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 }}
+1 -1
View File
@@ -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: |
+1 -1
View File
@@ -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: |
+2
View File
@@ -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
+2 -2
View File
@@ -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).
+3 -3
View File
@@ -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.11.0+cpu torchvision==0.26.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'])" && \
\
+2 -2
View File
@@ -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
+10 -1
View File
@@ -33,7 +33,16 @@ Path to save and load the embeddings index. Each API instance can only access a
writable: boolean
```
Determines if the input embeddings index is writable (true) or read-only (false). This allows serving a read-only index.
Sets this embeddings index as writable (true) or read-only (false). This allows serving a read-only index. Defaults to read-only.
### reindex
```yaml
reindex: boolean
```
Enables re-indexing for this embeddings index (defaults to disabled).
Re-indexing accepts new configuration and should only be enabled when the configuration comes from trusted sources.
### cloud
[Cloud storage settings](../../embeddings/configuration/cloud) can be set under a `cloud` top level configuration group.
+43 -2
View File
@@ -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
View File
@@ -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 .
+54 -7
View File
@@ -4,7 +4,7 @@ Approximate Nearest Neighbor (ANN) index configuration for storing vector embedd
## backend
```yaml
backend: faiss|hnsw|annoy|ggml|numpy|torch|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,6 +122,30 @@ torch:
blocksize: quantization block size parameter
```
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._
### 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
pgvector:
+2
View File
@@ -79,3 +79,5 @@ format: json|pickle
```
Sets the configuration storage format. Defaults to `json`.
Note that `pickle` configuration is deprecated and will raise an error with the default settings. Its only kept for reading legacy indexes and requires setting the `ALLOW_PICKLE` environment variable. Only enable for local and/or trusted sources.
+21 -2
View File
@@ -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
+43 -4
View File
@@ -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
+2
View File
@@ -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
View File
@@ -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
View File
@@ -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
+6
View File
@@ -17,6 +17,12 @@ Apache Tika requires [Java](https://en.wikipedia.org/wiki/Java_(programming_lang
See [this link](https://github.com/DS4SD/docling?tab=readme-ov-file#features) for a list of supported document formats.
## LiteParse
[LiteParse](https://github.com/run-llama/liteparse) is a fast, helpful, and open-source document parser. It supports PDF out of the box with additional format support via third-party packages.
See [this link](https://github.com/run-llama/liteparse#multi-format-input-support) for more.
## Example
The following shows a simple example using this pipeline.
+98
View File
@@ -0,0 +1,98 @@
# LemurTrainer
![pipeline](../../images/pipeline.png#only-light)
![pipeline](../../images/pipeline-dark.png#only-dark)
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__
+3 -2
View File
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
```yaml
workflow:
tasks:
- task: console
name:
tasks:
- task: console
```
## Methods
+3 -2
View File
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
```yaml
workflow:
tasks:
- task: export
name:
tasks:
- task: export
```
## Methods
+3 -2
View File
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
```yaml
workflow:
tasks:
- task: file
name:
tasks:
- task: file
```
## Methods
+3 -2
View File
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
```yaml
workflow:
tasks:
- task: image
name:
tasks:
- task: image
```
## Methods
+17 -3
View File
@@ -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],
+4 -3
View File
@@ -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
+4 -3
View File
@@ -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
+3 -2
View File
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
```yaml
workflow:
tasks:
- task: storage
name:
tasks:
- task: storage
```
## Methods
+4 -3
View File
@@ -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
+3 -2
View File
@@ -22,8 +22,9 @@ This task can also be created with workflow configuration.
```yaml
workflow:
tasks:
- task: url
name:
tasks:
- task: url
```
## Methods
+1 -1
View File
@@ -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"
]
}
],
+2 -2
View File
@@ -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",
+15 -15
View File
@@ -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",
+1
View File
@@ -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:
+21 -7
View File
@@ -20,7 +20,8 @@ default = [
"pyyaml>=5.3",
"safetensors>=0.4.5",
"torch>=2.4",
"transformers>=4.56.2",
"tqdm>=4.66.3",
"transformers>=5.9.0",
]
# Optional dependencies
@@ -52,18 +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-py>=0.9.4",
"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"] = [
@@ -96,13 +101,21 @@ extras["pipeline-audio"] = [
"webrtcvad-wheels>=2.0.14",
]
extras["pipeline-data"] = ["beautifulsoup4>=4.9.3", "chonkie>=1.0.2", "docling>=2.8.2", "nltk>=3.5", "pandas>=1.1.0", "tika>=1.24"]
extras["pipeline-data"] = [
"beautifulsoup4>=4.9.3",
"chonkie>=1.0.2",
"docling>=2.8.2",
"liteparse>=2.1.1",
"nltk>=3.5",
"pandas>=1.1.0",
"tika>=1.24",
]
extras["pipeline-image"] = ["imagehash>=4.2.1", "pillow>=7.1.2", "timm>=0.4.12"]
extras["pipeline-llm"] = ["httpx>=0.28.1", "litert-lm-api>=0.11.0", "litellm>=1.37.16", "llama-cpp-python>=0.2.75"]
extras["pipeline-text"] = ["gliner>=0.2.23", "sentencepiece>=0.1.91", "staticvectors>=0.2.0"]
extras["pipeline-text"] = ["gliner-py>=0.2.27", "sentencepiece>=0.1.91", "staticvectors>=0.2.0"]
extras["pipeline-train"] = [
"accelerate>=0.26.0",
@@ -130,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",
@@ -167,7 +181,7 @@ extras["all"] = (
setup(
name=package,
version="9.10.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,
@@ -182,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=[
+16 -4
View File
@@ -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
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", "find", "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
+26 -12
View File
@@ -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):
"""
+4
View File
@@ -6,6 +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
+17 -3
View File
@@ -4,14 +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:
@@ -31,6 +36,7 @@ class ANNFactory:
ANN
"""
# pylint: disable=R0912
# ANN instance
ann = None
backend = config.get("backend", "faiss" if FAISS else "numpy")
@@ -40,18 +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)
@@ -74,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
+2 -1
View File
@@ -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):
+36 -15
View File
@@ -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)
+5 -2
View File
@@ -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 = []
+144
View File
@@ -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()
+8 -4
View File
@@ -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)
+2 -2
View File
@@ -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 = []
+176
View File
@@ -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))
+21 -7
View File
@@ -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):
"""
+82
View File
@@ -0,0 +1,82 @@
"""
TurboVec module
"""
# Conditional import
try:
from turbovec import IdMapIndex
TURBOVEC = True
except ImportError:
TURBOVEC = False
from ..base import ANN
# Core library imports
from ...util import Library
np = Library().numpy()
class TurboVec(ANN):
"""
Builds an ANN index using the turbovec library.
"""
def __init__(self, config):
super().__init__(config)
if not TURBOVEC:
raise ImportError('turbovec is not available - install "ann" extra to enable')
def load(self, path):
# Load index
self.backend = IdMapIndex.load(path)
def index(self, embeddings):
# Lookup index settings
bitwidth = self.setting("bitwidth", 4)
# Create index
self.backend = IdMapIndex(dim=self.config["dimensions"], bit_width=bitwidth)
# Add items - position in embeddings is used as the id
self.backend.add_with_ids(embeddings, np.arange(embeddings.shape[0], dtype=np.uint64))
# Add id offset, delete counter and index build metadata
self.config["offset"] = embeddings.shape[0]
self.metadata({"bitwidth": bitwidth})
def append(self, embeddings):
new = embeddings.shape[0]
self.backend.add_with_ids(embeddings, np.arange(self.config["offset"], self.config["offset"] + new, dtype=np.uint64))
# Update id offset and index metadata
self.config["offset"] += new
self.metadata()
def delete(self, ids):
# Mark elements as deleted to omit from search results
for uid in ids:
self.backend.remove(uid)
def search(self, queries, limit):
# Run the query
scores, ids = self.backend.search(queries=queries, k=limit)
# Map results to [(id, score)]
results = []
for x, score in enumerate(scores):
# Build (id, score) tuples, convert to ints and floats
results.append(list(zip(ids[x].tolist(), score.tolist())))
return results
def count(self):
return len(self.backend)
def save(self, path):
# Write index
self.backend.write(path)
+166
View File
@@ -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()
+1
View File
@@ -5,3 +5,4 @@ Sparse ANN imports
from .factory import SparseANNFactory
from .ivfsparse import IVFSparse
from .pgsparse import PGSparse
from .zvecsparse import ZvecSparse
+5 -1
View File
@@ -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
+14 -6
View File
@@ -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):
+35
View File
@@ -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)}
+20 -4
View File
@@ -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.
+11 -14
View File
@@ -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()
+15 -1
View File
@@ -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)
+14 -12
View File
@@ -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"&parameters={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)
+9 -6
View File
@@ -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)
+24 -17
View File
@@ -2,9 +2,12 @@
Route module
"""
from fastapi.routing import APIRoute, get_request_handler
import json
from .responses import ResponseFactory
from fastapi.responses import StreamingResponse
from fastapi.routing import APIRoute
from .responses import ResponseFactory, MessagePackResponse
class EncodingAPIRoute(APIRoute):
@@ -20,22 +23,26 @@ class EncodingAPIRoute(APIRoute):
route handler function
"""
# Get handle to the original route handler
original = super().get_route_handler()
async def handler(request):
route = get_request_handler(
dependant=self.dependant,
body_field=self.body_field,
status_code=self.status_code,
response_class=ResponseFactory.create(request),
response_field=getattr(self, "secure_cloned_response_field", self.response_field),
response_model_include=self.response_model_include,
response_model_exclude=self.response_model_exclude,
response_model_by_alias=self.response_model_by_alias,
response_model_exclude_unset=self.response_model_exclude_unset,
response_model_exclude_defaults=self.response_model_exclude_defaults,
response_model_exclude_none=self.response_model_exclude_none,
dependency_overrides_provider=self.dependency_overrides_provider,
# Target response class
target = ResponseFactory.create(request)
request.state.response_class = target
# Get response
response = await original(request)
# Force MessagePackResponse when it's requested and response type doesn't match
return (
target(
content=json.loads(response.body),
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, StreamingResponse)) and target == MessagePackResponse
else response
)
return await route(request)
return handler
+2 -2
View File
@@ -14,7 +14,7 @@ from .. import application
from ..responses import ResponseFactory
from ..route import EncodingAPIRoute
from ...app import ReadOnlyError
from ...app import ReadOnlyError, ReindexDisabledError
from ...graph import Graph
router = APIRouter(route_class=EncodingAPIRoute)
@@ -198,7 +198,7 @@ def reindex(config: dict = Body(...), function: str = Body(default=None)):
try:
application.get().reindex(config, function)
except ReadOnlyError as e:
except (ReadOnlyError, ReindexDisabledError) as e:
raise HTTPException(status_code=403, detail=e.args[0]) from e
+36 -13
View File
@@ -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"
+1 -1
View File
@@ -2,4 +2,4 @@
App imports
"""
from .base import Application, ReadOnlyError
from .base import Application, ReadOnlyError, ReindexDisabledError
+23 -3
View File
@@ -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):
"""
@@ -566,6 +576,10 @@ class Application:
if not self.config.get("writable"):
raise ReadOnlyError("Attempting to reindex a read-only index (writable != True)")
# Raise error if not re-indexable
if not self.config.get("reindex"):
raise ReindexDisabledError("Attempting to reindex with reindexing disabled (reindex != True)")
if self.embeddings:
with self.lock:
# Resolve function, if necessary
@@ -826,3 +840,9 @@ class ReadOnlyError(Exception):
"""
Error raised when trying to modify a read-only index
"""
class ReindexDisabledError(Exception):
"""
Error raised when trying to reindex with reindexing disabled
"""
+1 -1
View File
@@ -46,6 +46,6 @@ class Compress:
directory = os.path.abspath(directory)
path = os.path.abspath(path)
prefix = os.path.commonprefix([directory, path])
prefix = os.path.commonpath([directory, path])
return prefix == directory
+2 -1
View File
@@ -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
+5
View File
@@ -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
+13 -7
View File
@@ -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"):
+13 -7
View File
@@ -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
+1 -1
View File
@@ -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])))
+36 -12
View File
@@ -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
+3
View File
@@ -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
+1 -1
View File
@@ -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):
+2 -1
View File
@@ -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
+3 -3
View File
@@ -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)
+2 -1
View File
@@ -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
+42 -13
View File
@@ -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
+10 -8
View File
@@ -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
"""
+9
View File
@@ -0,0 +1,9 @@
"""
Error module
"""
class SQLError(Exception):
"""
Raised for errors generated by user SQL queries
"""
+29 -3
View File
@@ -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
+5 -2
View File
@@ -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):
+8 -4
View File
@@ -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):
+3 -10
View File
@@ -6,8 +6,7 @@ from zipfile import BadZipFile
# Conditionally import dimensionality reduction libraries as they aren't installed by default
try:
# pylint: disable=W0611
import skops
import skops.io as sio
from sklearn.decomposition import TruncatedSVD
@@ -89,10 +88,7 @@ class Reducer:
# Dimensionality reduction
try:
# pylint: disable=C0415
from skops.io import load
self.model = load(path)
self.model = sio.load(path)
except (BadZipFile, KeyError):
# Backwards compatible support for pickled models
self.model = SerializeFactory.create("pickle").load(path)
@@ -105,7 +101,4 @@ class Reducer:
path: directory path to save model
"""
# pylint: disable=C0415
from skops.io import dump
dump(self.model, path)
sio.dump(self.model, 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:
+40 -10
View File
@@ -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])]
+2 -1
View File
@@ -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
+3 -3
View File
@@ -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
+4 -2
View File
@@ -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
+37 -16
View File
@@ -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
+2 -1
View File
@@ -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
+1 -2
View File
@@ -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"])
+9 -1
View File
@@ -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,
}
+4
View File
@@ -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()
+4 -2
View File
@@ -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
+3 -1
View File
@@ -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
+16 -3
View File
@@ -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
+242 -28
View File
@@ -3,6 +3,7 @@ Late module
"""
from .base import Pooling
from .lemur import Lemur
from .muvera import Muvera
# Core library imports
@@ -12,6 +13,7 @@ library = Library()
np = library.numpy()
safetensors = library.safetensors()
torch = library.torch()
Module = library.module()
class LatePooling(Pooling):
@@ -22,33 +24,36 @@ 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)
# Get linear weights path
config = self.load(path, "1_Dense/config.json")
# 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:
# PyLate weights format
name = "1_Dense/model.safetensors"
models = [f"{x['path']}/model.safetensors" for x in config if x["path"].endswith("_Dense")]
else:
# Stanford weights format
name = "model.safetensors"
models = ["model.safetensors"]
# Read model settings
self.qprefix, self.qlength, self.dprefix, self.dlength = self.settings(path, config)
# Load linear layer
path = Download()(path, name)
with safetensors.safe_open(filename=path, framework="pt") as f:
weights = f.get_tensor("linear.weight")
# Load linear model
self.linear = self.loadlinear(path, models)
# Load weights into linear layer
self.linear = torch.nn.Linear(weights.shape[1], weights.shape[0], bias=False, device=self.device, dtype=weights.dtype)
with torch.no_grad():
self.linear.weight.copy_(weights)
# Configure token centering
self.center = self.centersettings(center, self.linear, centerconfigured)
self.batches = None
def forward(self, **inputs):
"""
@@ -61,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)
@@ -76,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
@@ -104,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):
"""
@@ -142,4 +291,69 @@ 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):
"""
Loads linear model.
Args:
path: model path
models: list of paths to each dense model
Returns:
linear model
"""
# Load dense layers as a sequential model
layers = []
for model in models:
model = Download()(path, model)
with safetensors.safe_open(filename=model, framework="pt") as f:
dense = []
for name in ["linear.weight", "residual.weight"]:
if name in f.keys():
weights = f.get_tensor(name)
# Load weights into linear layer
model = torch.nn.Linear(weights.shape[1], weights.shape[0], bias=False, device=self.device, dtype=weights.dtype)
with torch.no_grad():
model.weight.copy_(weights)
dense.append(model)
layers.append(Dense(*dense))
return torch.nn.Sequential(*layers)
class Dense(Module):
"""
Dense layer. Supports multiple linear layers that sum into a final answer.
"""
def __init__(self, *modules):
"""
Create a Dense layer.
Args:
modules: list of modules to sum outputs
"""
super().__init__()
self.layers = torch.nn.ModuleList(modules)
def forward(self, x):
"""
Sums the outputs of each module for the input.
Args:
x: input
Returns:
sum of the outputs of each layer
"""
# Compute sum of all module outputs
return sum(layer(x) for layer in self.layers)
+412
View File
@@ -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)}")
+36
View File
@@ -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

Some files were not shown because too many files have changed in this diff Show More