bm25_pt - GitLite

KO | EN

gitlite — search

#python #hacktoberfest #llm #rust #linux #macos #nodejs #cli #kubernetes #vue #awesome-list #pytorch #mcp #open-source

bm25_pt

(★ 105)

minimal pytorch implementation of bm25 (with sparse tensors)

파일 탐색기

최종 버전 다운로드 (.zip)

- - test.yml
- __init__.py
- bm25.py
- basic_test.py
- real_test.py
- scrolls_test_corpus.p
- scrolls_test_queries.p
.gitignore
LICENSE.txt
README.md
requirements.txt
setup.py

# 설치 가이드

1. 코드 내려받기

git clone https://github.com/jxmorris12/bm25_pt

깃허브에서 프로젝트 코드 전체를 내 컴퓨터로 내려받습니다.

cd bm25_pt

방금 내려받은 프로젝트 폴더 안으로 이동합니다.

2. 공식 설치 스크립트

쉬움 추천

사전 준비물

Python 3 pip 명령어를 쓰려면 Python이 필요합니다.

pip install bm25_pt

PyPI에 배포된 패키지를 바로 설치합니다. 소스 클론이 필요 없습니다.

✅ 설치 후 새 터미널을 열고, 프로그램의 버전 확인 명령(예: --version)으로 정상 설치됐는지 확인하세요.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

3. Python

쉬움

사전 준비물

Git GitHub에서 프로젝트 코드를 내려받으려면 필요합니다.
Python 3 설치 시 'Add Python to PATH' 옵션을 꼭 체크하세요 (Windows).

pip install bm25_pt

PyPI에 배포된 패키지를 바로 설치합니다. 소스 클론이 필요 없습니다.

✅ 에러 메시지 없이 실행되고 터미널에 안내 문구가 출력되면 정상입니다.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# bm25-pt

A minimal implementation of [Okapi BM25](https://en.wikipedia.org/wiki/Okapi_BM25) using PyTorch. (Also uses [HuggingFace tokenizers](https://huggingface.co/docs/tokenizers/en/index) behind the scenes to tokenize text.)

```bash
pip install bm25_pt
```

## Usage

```python
from bm25_pt import BM25

bm25 = BM25()
corpus = [
    "A high weight in tf–idf is reached by a high term frequency",
    "(in the given document) and a low document frequency of the term",
    "in the whole collection of documents; the weights hence tend to filter",
    "out common terms. Since the ratio inside the idf's log function is always",
    "greater than or equal to 1, the value of idf (and tf–idf) is greater than or equal",
    "to 0. As a term appears in more documents, the ratio inside the logarithm approaches",
    "1, bringing the idf and tf–idf closer to 0.",
]
bm25.index(corpus)

queries = ["weights", "ratio logarithm"]
doc_scores = bm25.score_batch(queries)
print(doc_scores)
>> tensor([[0.0000, 0.0000, 1.4238, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 1.5317, 0.0000, 2.0203, 0.0000]])
```

can also call `score()` with a single example:

```python
doc_scores = bm25.score(queries[0])
print(doc_scores)
>> tensor([0.0000, 0.0000, 1.4238, 0.0000, 0.0000, 0.0000, 0.0000])
```
### Use your own tokenizer

You can use your own tokenizer if you want. Simply provide your tokenizer to the `BM25` constructor:

```python
from bm25_pt import BM25
import transformers

tokenizer = transformers.AutoTokenizer.from_pretrained("t5-base")
bm25 = BM25(tokenizer=tokenizer)
```

### Use a GPU

To perform operations on a GPU, it's easy, just pass a device to initialization:

```python
from bm25_pt import BM25

bm25 = BM25(device='cuda')
```

### Benchmarking on LoCO

I benchmarked this library on the [Stanford LoCo benchmark](https://hazyresearch.stanford.edu/blog/2024-01-11-m2-bert-retrieval). I'm not using any special tokenizer, so it's just using `bert-base-uncased` and matching based on subword tokens. We could probably get better performance by using a word-based model, perhaps by using a cased model, and certainly by tuning hyperparameters, especially `b` and `k1`.

| Retrieval Encoders  | Seq Len  | Tau Scrolls – Summ. | Tau Scrolls – Gov. | Tau Scrolls QMSUM | QASPER - Title | QASPER - Abstract | Average |
|---------------------|----------|---------------------|--------------------|-------------------|----------------|-------------------|---------|
|    bm25_pt   | infinity |         89.5        |        96.1        |        56.2       |      94.3      |        99.4       |   87.1  |
|  M2-BERT-2048 (80M) |   2048   |         81.8        |        94.7        |        58.5       |      87.3      |        95.5       |   83.6  |
|  M2-BERT-8192 (80M) |   8192   |         94.7        |        96.5        |        64.1       |      86.8      |        97.5       |   85.9  |
| M2-BERT-32768 (80M) |   32768  |         98.6        |        98.5        |        69.5       |      97.4      |        98.7       |   92.5  |

As you can see below, this library is very fast, even on CPU:

| Dataset | Size | Index time (s) | Scoring time (s) |
|---------------------|----------|---------------------|--------------------|
| Tau Scrolls – Summ. | 3673 | 43.2 | 14.5 |
| Tau Scrolls – Gov. | 17,457 | 622.7 | 328.2 |
| Tau Scrolls QMSUM | 272 | 5.5 | 0.2 |
| QASPER - Title | 416 | 6.1 | 0.2 |
| QASPER - Abstract | 416 | 5.9 | 0.2

I think that most of the time is spent tokenizing long documents. It could likely be sped up with a faster tokenizer and/or multiprocessing.

### Citation

```
@misc{morris2024bm25pt,
  author = {Morris, John X},
  title = {bm25-pt: A minimal implementation of Okapi BM25 using PyTorch},
  url = {https://github.com/jxmorris12/bm25-pt},
  year = {2024}
}
```