Indexing a large .jsonl that cannot be loaded in RAM

[liyucheng09]

@xhluca Hi, thanks for the nice project.

I have a large .jsonl file where each line got a text and an id.
My question is, as I cannot load all lines in the RAM, i cannot obtain all ids before I start .index().

But your implementation seems to require all ids for retriever = bm25s.BM25(corpus=corpus_id) before the indexing started.
Any ideas to handle this issue?

[xhluca]

The corpus you are passing to BM25 is optional, it is only used during retrieval and when you are saving/loading the model. You can just initialize without the corpus:

retriever = bm25s.BM25()

That said, if you still want to pass the jsonl file to the retriever's initialization call, but don't have enough memory to load it in memory, you can take a look at the bm25s.utils.corpus.JsonlCorpus class which lets you read a jsonl dynamically through memory mapping:

bm25s/bm25s/utils/corpus.py

Lines 101 to 184 in 73c7dea

	class JsonlCorpus:
	"""
	A class to read a jsonl file line by line using mmap, allowing extremely fast
	access to any line in the file. For example, you could access the N-th line
	of a 10GB file in a fraction of a second, returning a dictionary.
	Example
	--------
	Traditioanally, you would read a jsonl file line by line like this:
	```python
	import json
	data = [json.loads(line) for line in open("file.jsonl")]
	print(corpus[1000])
	```
	This is memory inefficient and has a large overhead. Instead, you can use this class:
	```python
	corpus = JsonlCorpus("file.jsonl")
	print(corpus[1000])
	```
	Which only loads the line you need into memory, and is much faster.
	"""
	def __init__(self, path, show_progress=True, leave_progress=True, save_index=True):
	self.path = path
	# if the index file does not exist, create it
	if os.path.exists(change_extension(path, ".mmindex.json")):
	self.mmindex = load_mmindex(path)
	else:
	logging.info("Creating index file for jsonl corpus")
	mmindex = find_newline_positions(
	path, show_progress=show_progress, leave_progress=leave_progress
	)
	if save_index:
	save_mmindex(mmindex, path)
	self.mmindex = mmindex
	self.file_obj = open(path, "r")
	self.mmap_obj = mmap.mmap(self.file_obj.fileno(), 0, access=mmap.ACCESS_READ)
	logging.info("Opened file and mmap objects")
	def __len__(self):
	return len(self.mmindex)
	def __getitem__(self, index):
	# handle multiple indices
	if isinstance(index, int):
	return json.loads(
	get_line(
	self.path,
	index,
	self.mmindex,
	file_obj=self.file_obj,
	mmap_obj=self.mmap_obj,
	)
	)
	if isinstance(index, slice):
	return [self.__getitem__(i) for i in range(*index.indices(len(self)))]
	if isinstance(index, (list, tuple)):
	return [self.__getitem__(i) for i in index]
	if isinstance(index, np.ndarray):
	# if it's an ndarray, this means each element is an index, and the array can
	# be of any shape. thus, we should flatten it first, get the results as if it
	# was a list, and then reshape it back to the original shape
	index_flat = index.flatten().tolist()
	results = [self.__getitem__(i) for i in index_flat]
	reshaped = np.array(results).reshape(index.shape)
	return reshaped
	raise TypeError("Invalid index type")
	def __del__(self):
	if hasattr(self, "file_obj"):
	self.file_obj.close()
	if hasattr(self, "mmap_obj"):
	self.mmap_obj.close()
	logging.info("Closed file and mmap objects")

Note however that before it does that, it needs to index the jsonl and create a mmindex file in the same directory as the jsonl file; after that, it will read any line of the jsonl extremely quickly, and even has basic support for slicing!

[ddofer]

Any tip for doing this over multiple jsonl files? (That add up to a large one)? (For doing indexing over them all)?

[liyucheng09]

@xhluca This is really cool. And thanks for your prompt response!

One quick question, I have tried to new a BM25 without a given corpus, and I found results, _ = retriever.retrieve will just return the index of the document in the corpus_tokens.

I mean, it will use default ids assigned by the BM25 if the corpus is not given for it? Is there a way to pass the id along side the tokenized corpus? like retriever.index(corpus_tokens) where corpus_tokens are a list of dict [...{'id': ..., 'text': ...}...]. So that the final result can be my customized ids.

[xhluca]

Yes, it will return the index when the corpus is not available. If you want it to return ids, you can use retriever.retrieve(..., corpus=corpus_ids), or use the JsonlCorpus class I recommend above, e.g. try:

corpus = JsonlCorpus('corpus_records.jsonl')
lines = corpus[indexes_returned_by_retriever]

The reason index() does not store the corpus (or any key in the corpus) is that it is unpredictable how much memory will be used the iterator; as an extension, storing certain keys but not others would add complexity to the index API (which already has many params) without clear benefit over using JsonlCorpus or passing a corpus_ids to retriever.retrieve(corpus=corpus_ids).

If you still believe it is the best thing to do, feel free to subclass index:

class CustomBM25(bm25s.BM25):
  def index(self, ...):
    # modify the index implementation here

[xhluca]

You could also read the jsonl yourself and create a cached ids file, i.e.:

ids = []

with open('corpus.jsonl') as f:
  for line in f:
    record = json.load(line)
    ids.append(record['id'])

with open('corpus.ids.json', 'w') as f:
  json.dump(ids, f)