BM25 0.3.0rc5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,173 @@
1
+ Metadata-Version: 2.4
2
+ Name: BM25
3
+ Version: 0.3.0rc5
4
+ Summary: A simple high-level API and CLI for BM25.
5
+ Home-page: https://github.com/xhluca/bm25s/tree/main/bm25s/high_level
6
+ Author: Xing Han Lù
7
+ Author-email: bm25s@googlegroups.com
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.8
12
+ Description-Content-Type: text/markdown
13
+ Requires-Dist: bm25s[cli,core]==0.3.0.rc5
14
+ Dynamic: author
15
+ Dynamic: author-email
16
+ Dynamic: classifier
17
+ Dynamic: description
18
+ Dynamic: description-content-type
19
+ Dynamic: home-page
20
+ Dynamic: requires-dist
21
+ Dynamic: requires-python
22
+ Dynamic: summary
23
+
24
+ <div align="center">
25
+
26
+ <h1>BM25</h1>
27
+
28
+ <i>A fast, simple, and high-level Python API and CLI for BM25, powered by `bm25s`.</i>
29
+
30
+ <table>
31
+ <tr>
32
+ <td>
33
+ <a href="https://github.com/xhluca/bm25s">💻 GitHub</a>
34
+ </td>
35
+ <td>
36
+ <a href="https://pypi.org/project/bm25s/">📦 bm25s</a>
37
+ </td>
38
+ <td>
39
+ <a href="https://bm25s.github.io">🏠 Homepage</a>
40
+ </td>
41
+ </tr>
42
+ </table>
43
+ </div>
44
+
45
+ `BM25` is a wrapper package that installs `bm25s` with its optional core dependencies, providing a simple, high-level API and a command-line interface for fast and effective text retrieval.
46
+
47
+ ## Installation
48
+
49
+ Install `BM25` using pip:
50
+
51
+ ```bash
52
+ pip install BM25
53
+ ```
54
+
55
+ This will automatically install the highly optimized `bm25s` backend, alongside necessary dependencies for stemming (`PyStemmer`), parallelization, and CLI (`rich`).
56
+
57
+ ## High Level API
58
+
59
+ If you want to quickly search on a local file, you can use the `BM25` module:
60
+
61
+ ```python
62
+ import BM25
63
+
64
+ # Load a file (csv, json, jsonl, txt)
65
+ # For csv/jsonl, you can specify the column/key to use as document text
66
+ corpus = BM25.load("tests/data/dummy.csv", document_column="text")
67
+ # Index the corpus
68
+ retriever = BM25.index(corpus)
69
+
70
+ # Search
71
+ results = retriever.search(["your query here"], k=5)
72
+ for result in results[0]:
73
+ print(result)
74
+ ```
75
+
76
+ The `load` function handles file reading, while `index` handles tokenization, indexing, and provides a simple search interface.
77
+
78
+ ## Command-Line Interface
79
+
80
+ The package provides a terminal-based CLI for quick indexing and searching without writing Python code.
81
+
82
+ ### Indexing Documents
83
+
84
+ Create an index from a CSV, TXT, JSON, or JSONL file:
85
+
86
+ ```bash
87
+ # Index a CSV file (uses first column by default)
88
+ bm25 index documents.csv -o my_index
89
+
90
+ # Index with a specific column
91
+ bm25 index documents.csv -o my_index -c text
92
+
93
+ # Index a text file (one document per line)
94
+ bm25 index documents.txt -o my_index
95
+
96
+ # Index a JSONL file
97
+ bm25 index documents.jsonl -o my_index -c content
98
+ ```
99
+
100
+ If you don't specify an output directory with `-o`, the index will be saved to `<filename>_index`.
101
+
102
+ ### User Directory
103
+
104
+ You can save indices to a central user directory (`~/.bm25s/indices/`) using the `-u` flag:
105
+
106
+ ```bash
107
+ # Save index to ~/.bm25s/indices/my_docs
108
+ bm25 index documents.csv -u -o my_docs
109
+
110
+ # Search using the user directory
111
+ bm25 search -u -i my_docs "your query"
112
+ ```
113
+
114
+ ### Searching
115
+
116
+ Search an existing index with a query using `-i` (or `--index`):
117
+
118
+ ```bash
119
+ # Basic search (returns top 10 results)
120
+ bm25 search -i my_index "what is machine learning?"
121
+
122
+ # Search with full path
123
+ bm25 search -i ./path/to/my_index "your query here"
124
+
125
+ # Return more results
126
+ bm25 search -i my_index "your query here" -k 20
127
+
128
+ # Save results to a JSON file
129
+ bm25 search -i my_index "your query here" -s results.json
130
+ ```
131
+
132
+ ### Interactive Index Picker
133
+
134
+ When using `-u` without specifying an index name, an interactive picker is displayed (requires `bm25s[cli]` which is installed by default with `BM25`):
135
+
136
+ ```bash
137
+ # Interactive picker will show available indices
138
+ bm25 search -u "your query"
139
+ ```
140
+
141
+ ### Example Workflow
142
+
143
+ **Basic usage** (index saved to current directory):
144
+
145
+ ```bash
146
+ # 1. Create a simple text file with documents
147
+ echo -e "Machine learning is a subset of AI\nDeep learning uses neural networks\nNatural language processing handles text" > docs.txt
148
+
149
+ # 2. Index the documents
150
+ bm25 index docs.txt -o my_index
151
+
152
+ # 3. Search the index
153
+ bm25 search -i my_index "what is AI?"
154
+ ```
155
+
156
+ **With user directory** (indices saved to `~/.bm25s/indices/`):
157
+
158
+ ```bash
159
+ # Index to user directory
160
+ bm25 index docs.txt -u -o ml_docs
161
+
162
+ # Search from user directory
163
+ bm25 search -u -i ml_docs "what is AI?"
164
+
165
+ # Or use the interactive picker
166
+ bm25 search -u "what is AI?"
167
+ ```
168
+
169
+ ## Flexibility
170
+
171
+ For more advanced use cases, including memory mapping, customized tokenization, hugging face integration, or using different BM25 variants, please use the underlying `bm25s` API directly.
172
+
173
+ See the [bm25s documentation](https://github.com/xhluca/bm25s) for full details.
@@ -0,0 +1,9 @@
1
+ README.md
2
+ setup.py
3
+ ./__init__.py
4
+ BM25.egg-info/PKG-INFO
5
+ BM25.egg-info/SOURCES.txt
6
+ BM25.egg-info/dependency_links.txt
7
+ BM25.egg-info/entry_points.txt
8
+ BM25.egg-info/requires.txt
9
+ BM25.egg-info/top_level.txt
@@ -0,0 +1,2 @@
1
+ [console_scripts]
2
+ bm25 = bm25s.cli:main
@@ -0,0 +1 @@
1
+ bm25s[cli,core]==0.3.0.rc5
@@ -0,0 +1 @@
1
+ BM25
bm25-0.3.0rc5/PKG-INFO ADDED
@@ -0,0 +1,173 @@
1
+ Metadata-Version: 2.4
2
+ Name: BM25
3
+ Version: 0.3.0rc5
4
+ Summary: A simple high-level API and CLI for BM25.
5
+ Home-page: https://github.com/xhluca/bm25s/tree/main/bm25s/high_level
6
+ Author: Xing Han Lù
7
+ Author-email: bm25s@googlegroups.com
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.8
12
+ Description-Content-Type: text/markdown
13
+ Requires-Dist: bm25s[cli,core]==0.3.0.rc5
14
+ Dynamic: author
15
+ Dynamic: author-email
16
+ Dynamic: classifier
17
+ Dynamic: description
18
+ Dynamic: description-content-type
19
+ Dynamic: home-page
20
+ Dynamic: requires-dist
21
+ Dynamic: requires-python
22
+ Dynamic: summary
23
+
24
+ <div align="center">
25
+
26
+ <h1>BM25</h1>
27
+
28
+ <i>A fast, simple, and high-level Python API and CLI for BM25, powered by `bm25s`.</i>
29
+
30
+ <table>
31
+ <tr>
32
+ <td>
33
+ <a href="https://github.com/xhluca/bm25s">💻 GitHub</a>
34
+ </td>
35
+ <td>
36
+ <a href="https://pypi.org/project/bm25s/">📦 bm25s</a>
37
+ </td>
38
+ <td>
39
+ <a href="https://bm25s.github.io">🏠 Homepage</a>
40
+ </td>
41
+ </tr>
42
+ </table>
43
+ </div>
44
+
45
+ `BM25` is a wrapper package that installs `bm25s` with its optional core dependencies, providing a simple, high-level API and a command-line interface for fast and effective text retrieval.
46
+
47
+ ## Installation
48
+
49
+ Install `BM25` using pip:
50
+
51
+ ```bash
52
+ pip install BM25
53
+ ```
54
+
55
+ This will automatically install the highly optimized `bm25s` backend, alongside necessary dependencies for stemming (`PyStemmer`), parallelization, and CLI (`rich`).
56
+
57
+ ## High Level API
58
+
59
+ If you want to quickly search on a local file, you can use the `BM25` module:
60
+
61
+ ```python
62
+ import BM25
63
+
64
+ # Load a file (csv, json, jsonl, txt)
65
+ # For csv/jsonl, you can specify the column/key to use as document text
66
+ corpus = BM25.load("tests/data/dummy.csv", document_column="text")
67
+ # Index the corpus
68
+ retriever = BM25.index(corpus)
69
+
70
+ # Search
71
+ results = retriever.search(["your query here"], k=5)
72
+ for result in results[0]:
73
+ print(result)
74
+ ```
75
+
76
+ The `load` function handles file reading, while `index` handles tokenization, indexing, and provides a simple search interface.
77
+
78
+ ## Command-Line Interface
79
+
80
+ The package provides a terminal-based CLI for quick indexing and searching without writing Python code.
81
+
82
+ ### Indexing Documents
83
+
84
+ Create an index from a CSV, TXT, JSON, or JSONL file:
85
+
86
+ ```bash
87
+ # Index a CSV file (uses first column by default)
88
+ bm25 index documents.csv -o my_index
89
+
90
+ # Index with a specific column
91
+ bm25 index documents.csv -o my_index -c text
92
+
93
+ # Index a text file (one document per line)
94
+ bm25 index documents.txt -o my_index
95
+
96
+ # Index a JSONL file
97
+ bm25 index documents.jsonl -o my_index -c content
98
+ ```
99
+
100
+ If you don't specify an output directory with `-o`, the index will be saved to `<filename>_index`.
101
+
102
+ ### User Directory
103
+
104
+ You can save indices to a central user directory (`~/.bm25s/indices/`) using the `-u` flag:
105
+
106
+ ```bash
107
+ # Save index to ~/.bm25s/indices/my_docs
108
+ bm25 index documents.csv -u -o my_docs
109
+
110
+ # Search using the user directory
111
+ bm25 search -u -i my_docs "your query"
112
+ ```
113
+
114
+ ### Searching
115
+
116
+ Search an existing index with a query using `-i` (or `--index`):
117
+
118
+ ```bash
119
+ # Basic search (returns top 10 results)
120
+ bm25 search -i my_index "what is machine learning?"
121
+
122
+ # Search with full path
123
+ bm25 search -i ./path/to/my_index "your query here"
124
+
125
+ # Return more results
126
+ bm25 search -i my_index "your query here" -k 20
127
+
128
+ # Save results to a JSON file
129
+ bm25 search -i my_index "your query here" -s results.json
130
+ ```
131
+
132
+ ### Interactive Index Picker
133
+
134
+ When using `-u` without specifying an index name, an interactive picker is displayed (requires `bm25s[cli]` which is installed by default with `BM25`):
135
+
136
+ ```bash
137
+ # Interactive picker will show available indices
138
+ bm25 search -u "your query"
139
+ ```
140
+
141
+ ### Example Workflow
142
+
143
+ **Basic usage** (index saved to current directory):
144
+
145
+ ```bash
146
+ # 1. Create a simple text file with documents
147
+ echo -e "Machine learning is a subset of AI\nDeep learning uses neural networks\nNatural language processing handles text" > docs.txt
148
+
149
+ # 2. Index the documents
150
+ bm25 index docs.txt -o my_index
151
+
152
+ # 3. Search the index
153
+ bm25 search -i my_index "what is AI?"
154
+ ```
155
+
156
+ **With user directory** (indices saved to `~/.bm25s/indices/`):
157
+
158
+ ```bash
159
+ # Index to user directory
160
+ bm25 index docs.txt -u -o ml_docs
161
+
162
+ # Search from user directory
163
+ bm25 search -u -i ml_docs "what is AI?"
164
+
165
+ # Or use the interactive picker
166
+ bm25 search -u "what is AI?"
167
+ ```
168
+
169
+ ## Flexibility
170
+
171
+ For more advanced use cases, including memory mapping, customized tokenization, hugging face integration, or using different BM25 variants, please use the underlying `bm25s` API directly.
172
+
173
+ See the [bm25s documentation](https://github.com/xhluca/bm25s) for full details.
@@ -0,0 +1,150 @@
1
+ <div align="center">
2
+
3
+ <h1>BM25</h1>
4
+
5
+ <i>A fast, simple, and high-level Python API and CLI for BM25, powered by `bm25s`.</i>
6
+
7
+ <table>
8
+ <tr>
9
+ <td>
10
+ <a href="https://github.com/xhluca/bm25s">💻 GitHub</a>
11
+ </td>
12
+ <td>
13
+ <a href="https://pypi.org/project/bm25s/">📦 bm25s</a>
14
+ </td>
15
+ <td>
16
+ <a href="https://bm25s.github.io">🏠 Homepage</a>
17
+ </td>
18
+ </tr>
19
+ </table>
20
+ </div>
21
+
22
+ `BM25` is a wrapper package that installs `bm25s` with its optional core dependencies, providing a simple, high-level API and a command-line interface for fast and effective text retrieval.
23
+
24
+ ## Installation
25
+
26
+ Install `BM25` using pip:
27
+
28
+ ```bash
29
+ pip install BM25
30
+ ```
31
+
32
+ This will automatically install the highly optimized `bm25s` backend, alongside necessary dependencies for stemming (`PyStemmer`), parallelization, and CLI (`rich`).
33
+
34
+ ## High Level API
35
+
36
+ If you want to quickly search on a local file, you can use the `BM25` module:
37
+
38
+ ```python
39
+ import BM25
40
+
41
+ # Load a file (csv, json, jsonl, txt)
42
+ # For csv/jsonl, you can specify the column/key to use as document text
43
+ corpus = BM25.load("tests/data/dummy.csv", document_column="text")
44
+ # Index the corpus
45
+ retriever = BM25.index(corpus)
46
+
47
+ # Search
48
+ results = retriever.search(["your query here"], k=5)
49
+ for result in results[0]:
50
+ print(result)
51
+ ```
52
+
53
+ The `load` function handles file reading, while `index` handles tokenization, indexing, and provides a simple search interface.
54
+
55
+ ## Command-Line Interface
56
+
57
+ The package provides a terminal-based CLI for quick indexing and searching without writing Python code.
58
+
59
+ ### Indexing Documents
60
+
61
+ Create an index from a CSV, TXT, JSON, or JSONL file:
62
+
63
+ ```bash
64
+ # Index a CSV file (uses first column by default)
65
+ bm25 index documents.csv -o my_index
66
+
67
+ # Index with a specific column
68
+ bm25 index documents.csv -o my_index -c text
69
+
70
+ # Index a text file (one document per line)
71
+ bm25 index documents.txt -o my_index
72
+
73
+ # Index a JSONL file
74
+ bm25 index documents.jsonl -o my_index -c content
75
+ ```
76
+
77
+ If you don't specify an output directory with `-o`, the index will be saved to `<filename>_index`.
78
+
79
+ ### User Directory
80
+
81
+ You can save indices to a central user directory (`~/.bm25s/indices/`) using the `-u` flag:
82
+
83
+ ```bash
84
+ # Save index to ~/.bm25s/indices/my_docs
85
+ bm25 index documents.csv -u -o my_docs
86
+
87
+ # Search using the user directory
88
+ bm25 search -u -i my_docs "your query"
89
+ ```
90
+
91
+ ### Searching
92
+
93
+ Search an existing index with a query using `-i` (or `--index`):
94
+
95
+ ```bash
96
+ # Basic search (returns top 10 results)
97
+ bm25 search -i my_index "what is machine learning?"
98
+
99
+ # Search with full path
100
+ bm25 search -i ./path/to/my_index "your query here"
101
+
102
+ # Return more results
103
+ bm25 search -i my_index "your query here" -k 20
104
+
105
+ # Save results to a JSON file
106
+ bm25 search -i my_index "your query here" -s results.json
107
+ ```
108
+
109
+ ### Interactive Index Picker
110
+
111
+ When using `-u` without specifying an index name, an interactive picker is displayed (requires `bm25s[cli]` which is installed by default with `BM25`):
112
+
113
+ ```bash
114
+ # Interactive picker will show available indices
115
+ bm25 search -u "your query"
116
+ ```
117
+
118
+ ### Example Workflow
119
+
120
+ **Basic usage** (index saved to current directory):
121
+
122
+ ```bash
123
+ # 1. Create a simple text file with documents
124
+ echo -e "Machine learning is a subset of AI\nDeep learning uses neural networks\nNatural language processing handles text" > docs.txt
125
+
126
+ # 2. Index the documents
127
+ bm25 index docs.txt -o my_index
128
+
129
+ # 3. Search the index
130
+ bm25 search -i my_index "what is AI?"
131
+ ```
132
+
133
+ **With user directory** (indices saved to `~/.bm25s/indices/`):
134
+
135
+ ```bash
136
+ # Index to user directory
137
+ bm25 index docs.txt -u -o ml_docs
138
+
139
+ # Search from user directory
140
+ bm25 search -u -i ml_docs "what is AI?"
141
+
142
+ # Or use the interactive picker
143
+ bm25 search -u "what is AI?"
144
+ ```
145
+
146
+ ## Flexibility
147
+
148
+ For more advanced use cases, including memory mapping, customized tokenization, hugging face integration, or using different BM25 variants, please use the underlying `bm25s` API directly.
149
+
150
+ See the [bm25s documentation](https://github.com/xhluca/bm25s) for full details.
@@ -0,0 +1,252 @@
1
+ """
2
+ The high level BM25 search API. It wraps bm25s into a simple to use search interface,
3
+ enabling 1-line indexing and 1-line searching. By default, it will require:
4
+ - numba compilation for speed up
5
+ - stemming for better search quality
6
+ - stopword removal for better search quality
7
+ """
8
+
9
+ import json
10
+ import csv
11
+ from pathlib import Path
12
+ from bm25s import BM25
13
+ from bm25s.tokenization import Tokenizer
14
+ import Stemmer
15
+ from typing import List
16
+
17
+
18
+
19
+ class BM25Search:
20
+ def __init__(
21
+ self,
22
+ corpus: List[str],
23
+ language: str = "english",
24
+ bm25_kwargs: dict = None,
25
+ tokenizer_kwargs: dict = None,
26
+ tokenizer_cls: Tokenizer = Tokenizer,
27
+ ):
28
+ if bm25_kwargs is not None and "corpus" in bm25_kwargs:
29
+ raise ValueError(
30
+ "The 'corpus' argument in bm25_kwargs is reserved and cannot be set manually."
31
+ )
32
+ if language != "english":
33
+ raise NotImplementedError("Currently only English language is supported.")
34
+
35
+ self.leave_progress = leave_progress = False
36
+ self.show_progress = show_progress = True
37
+ self.corpus = corpus
38
+
39
+ stemmer = Stemmer.Stemmer("english")
40
+ bm25_kwargs_default = dict(
41
+ backend="numba", csc_backend="numpy", auto_compile=False
42
+ )
43
+ tokenizer_kwargs_default = dict(
44
+ stemmer=stemmer, stopwords="english", lower=True
45
+ )
46
+
47
+ if isinstance(bm25_kwargs, dict):
48
+ bm25_kwargs_default.update(bm25_kwargs)
49
+ elif bm25_kwargs is not None:
50
+ raise ValueError("bm25_kwargs must be a dict or None.")
51
+
52
+ if isinstance(tokenizer_kwargs, dict):
53
+ tokenizer_kwargs_default.update(tokenizer_kwargs)
54
+ elif tokenizer_kwargs is not None:
55
+ raise ValueError("tokenizer_kwargs must be a dict or None.")
56
+
57
+ if "corpus" in bm25_kwargs_default:
58
+ raise ValueError(
59
+ "The 'corpus' argument in bm25_kwargs is reserved and cannot be set manually."
60
+ )
61
+
62
+ # note: we do not pass corpus here, as we will keep it separately. This means the BM25
63
+ # object will return document ids instead of texts when retrieving.
64
+ self.retriever = BM25(**bm25_kwargs_default)
65
+ self.tokenizer: Tokenizer = tokenizer_cls(**tokenizer_kwargs_default)
66
+
67
+ # tokenize the corpus
68
+ tokenized = self.tokenizer.tokenize(
69
+ corpus,
70
+ leave_progress=leave_progress,
71
+ show_progress=show_progress,
72
+ update_vocab=True,
73
+ return_as="tuple",
74
+ )
75
+
76
+ # compile and index (warmup=False to avoid segfaults in some CI environments)
77
+ self.retriever.compile(activate_numba=True, warmup=False)
78
+
79
+ create_empty_token = True
80
+ # If the corpus is empty or has no tokens, we can't create an empty token
81
+ # as it relies on vocab dict having some content or logic that fails if empty
82
+ if len(tokenized.vocab) == 0:
83
+ create_empty_token = False
84
+
85
+ self.retriever.index(
86
+ tokenized,
87
+ leave_progress=leave_progress,
88
+ show_progress=show_progress,
89
+ create_empty_token=create_empty_token,
90
+ )
91
+
92
+ def search(self, queries: List[str], k: int = 10, n_jobs: int = 1):
93
+ # Ensure k is not larger than the corpus size
94
+ num_docs = len(self.corpus)
95
+ if k > num_docs:
96
+ k = num_docs
97
+
98
+ tokenized_queries = self.tokenizer.tokenize(
99
+ queries,
100
+ update_vocab=False,
101
+ show_progress=self.show_progress,
102
+ leave_progress=self.leave_progress,
103
+ return_as="tuple",
104
+ allow_empty=False,
105
+ )
106
+
107
+ # Handle empty queries explicitly to avoid issues with Numba backend
108
+ # We filter out queries that result in no tokens
109
+ non_empty_indices = []
110
+ empty_indices = []
111
+
112
+ # We access internal ids list from Tokenized namedtuple
113
+ # Convert to string list first as retrieve expects that or Tokenized object
114
+ # But checking emptiness is easier on ids
115
+
116
+ for i, q_ids in enumerate(tokenized_queries.ids):
117
+ if len(q_ids) > 0:
118
+ non_empty_indices.append(i)
119
+ else:
120
+ empty_indices.append(i)
121
+
122
+ # Prepare results structure
123
+ results = [None] * len(queries)
124
+
125
+ # Process empty queries immediately
126
+ for i in empty_indices:
127
+ # Empty query results in empty list of documents
128
+ # Or do we want to return empty docs with 0 score?
129
+ # Standard BM25 usually returns nothing for empty query, or 0 score for all docs.
130
+ # Let's return empty list as top-k results.
131
+ results[i] = []
132
+
133
+ if len(non_empty_indices) > 0:
134
+ # Create a new Tokenized object for non-empty queries
135
+ non_empty_ids = [tokenized_queries.ids[i] for i in non_empty_indices]
136
+ non_empty_tokenized = self.tokenizer.to_tokenized_tuple(non_empty_ids)
137
+
138
+ # Retrieve for non-empty queries
139
+ # note: because we did not pass `corpus` when initializing BM25,
140
+ # the retrieve() will return document ids instead of texts.
141
+ doc_ids, scores = self.retriever.retrieve(
142
+ query_tokens=non_empty_tokenized,
143
+ k=k,
144
+ sorted=True,
145
+ return_as="tuple",
146
+ show_progress=self.show_progress,
147
+ leave_progress=self.leave_progress,
148
+ n_threads=n_jobs,
149
+ chunksize=50,
150
+ backend_selection="auto",
151
+ )
152
+
153
+ # Map back to original indices
154
+ num_docs = doc_ids.shape[1]
155
+
156
+ for idx, original_idx in enumerate(non_empty_indices):
157
+ query_results = []
158
+ for di in range(num_docs):
159
+ doc_id = doc_ids[idx, di]
160
+ doc_text = self.corpus[doc_id]
161
+ query_results.append(
162
+ {
163
+ "id": int(doc_id),
164
+ "score": float(scores[idx, di]),
165
+ "document": doc_text,
166
+ }
167
+ )
168
+ results[original_idx] = query_results
169
+
170
+ return results
171
+
172
+
173
+ def index(documents, language: str = "english"):
174
+ return BM25Search(corpus=documents, language=language)
175
+
176
+
177
+ def load(path, document_column=None):
178
+ """
179
+ Loads a csv, json, jsonl, or txt file from the given path. For json, we expect a list of dicts.
180
+ Returns a list of strings (corpus) that can be passed to `bm25s.index`.
181
+
182
+ Parameters
183
+ ----------
184
+ path : str
185
+ The file path to load the documents from.
186
+
187
+ document_column : str, optional
188
+ The column name to use as the document text when loading from csv, or the key when loading from json/jsonl.
189
+ If None, the first column will be used by default.
190
+
191
+ Returns
192
+ -------
193
+ List[str]
194
+ A list of strings representing the documents.
195
+ """
196
+
197
+ path = Path(path)
198
+ documents = []
199
+
200
+ if path.suffix == ".txt":
201
+ with open(path, "r", encoding="utf-8") as f:
202
+ documents = [line.strip() for line in f if line.strip()]
203
+
204
+ elif path.suffix == ".json":
205
+ with open(path, "r", encoding="utf-8") as f:
206
+ data = json.load(f)
207
+ if isinstance(data, list):
208
+ if len(data) > 0 and isinstance(data[0], str):
209
+ documents = data
210
+ elif len(data) > 0 and isinstance(data[0], dict):
211
+ if document_column is None:
212
+ # Use the first key available in the first element
213
+ document_column = list(data[0].keys())[0]
214
+ documents = [d[document_column] for d in data]
215
+ else:
216
+ raise ValueError("JSON file must contain a list of strings or dicts.")
217
+
218
+ elif path.suffix == ".jsonl":
219
+ with open(path, "r", encoding="utf-8") as f:
220
+ for line in f:
221
+ if not line.strip():
222
+ continue
223
+ d = json.loads(line)
224
+ if document_column is None and not documents:
225
+ # Infer column from first line
226
+ document_column = list(d.keys())[0]
227
+
228
+ if document_column in d:
229
+ documents.append(d[document_column])
230
+ else:
231
+ # skip or error? Let's skip if key missing, or error.
232
+ # raising error is safer for "load"
233
+ raise ValueError(f"Key '{document_column}' not found in JSONL line: {line[:50]}...")
234
+
235
+ elif path.suffix == ".csv":
236
+ with open(path, "r", newline="", encoding="utf-8") as f:
237
+ reader = csv.DictReader(f)
238
+ if document_column is None:
239
+ if reader.fieldnames:
240
+ document_column = reader.fieldnames[0]
241
+ else:
242
+ return BM25Search(corpus=[])
243
+
244
+ for row in reader:
245
+ if document_column in row:
246
+ documents.append(row[document_column])
247
+ else:
248
+ raise ValueError(f"Column '{document_column}' not found in CSV.")
249
+ else:
250
+ raise ValueError(f"Unsupported file extension: {path.suffix}")
251
+
252
+ return documents
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
bm25-0.3.0rc5/setup.py ADDED
@@ -0,0 +1,42 @@
1
+ from setuptools import setup
2
+ import os
3
+ import sys
4
+
5
+ # Change to the current directory to avoid issues if run from elsewhere
6
+ current_dir = os.path.dirname(os.path.abspath(__file__))
7
+ os.chdir(current_dir)
8
+
9
+ package_name = "BM25"
10
+ version = {}
11
+ with open(os.path.join("..", "version.py"), encoding="utf8") as fp:
12
+ exec(fp.read(), version)
13
+
14
+ with open("README.md", encoding="utf8") as fp:
15
+ long_description = fp.read()
16
+
17
+ setup(
18
+ name=package_name,
19
+ version=version["__version__"],
20
+ author="Xing Han Lù",
21
+ author_email="bm25s@googlegroups.com",
22
+ url="https://github.com/xhluca/bm25s/tree/main/bm25s/high_level",
23
+ description="A simple high-level API and CLI for BM25.",
24
+ long_description=long_description,
25
+ long_description_content_type="text/markdown",
26
+ packages=["BM25"],
27
+ package_dir={"BM25": "."},
28
+ install_requires=[
29
+ f"bm25s[core,cli]=={version['__version__']}",
30
+ ],
31
+ entry_points={
32
+ "console_scripts": [
33
+ "bm25=bm25s.cli:main",
34
+ ],
35
+ },
36
+ classifiers=[
37
+ "Programming Language :: Python :: 3",
38
+ "License :: OSI Approved :: MIT License",
39
+ "Operating System :: OS Independent",
40
+ ],
41
+ python_requires=">=3.8",
42
+ )