116 lines
4.3 KiB
Python
116 lines
4.3 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# pylint: disable=missing-module-docstring, missing-class-docstring
|
|
from __future__ import annotations
|
|
import typing
|
|
|
|
from flask_babel import gettext
|
|
|
|
from searx import settings
|
|
from searx.plugins import Plugin, PluginInfo
|
|
from searx.result_types import EngineResults
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from searx.search import SearchWithPlugins
|
|
from searx.extended_types import SXNG_Request
|
|
|
|
try:
|
|
import bm25s
|
|
except ImportError:
|
|
# Import error is ignored because the admin has to install bm25s manually to use the plugin
|
|
bm25s = None
|
|
|
|
|
|
class SXNGPlugin(Plugin):
|
|
"""Plugin which reranks the search results using the Okapi BM25 algorithm.
|
|
|
|
This plugin utilizes the `bm25s` library to reorder search results based on their relevance to the search query,
|
|
potentially improving the quality of results. Before enabling this plugin,
|
|
ensure you have installed the ``bm25s`` pip package. e.g. by installing it directly via pip or
|
|
by adding it to the project's `requirements.txt` file.
|
|
|
|
Configuration:
|
|
--------------
|
|
To enable the Rerank plugin, add it to the `enabled_plugins` list in your `settings.yml` file:
|
|
|
|
.. code:: yaml
|
|
|
|
enabled_plugins:
|
|
..
|
|
- 'Rerank plugin'
|
|
|
|
By default, the plugin retains the information about which engines found a particular result.
|
|
Results that appear in multiple engine results will receive a score boost.
|
|
This approach might be relevant if you wish results found by different engines to be prioritized.
|
|
You can modify this behaviour by configuring the ``remove_extra_engines`` setting.
|
|
If ``remove_extra_engines`` is set to ``true``, the original engine list is reduced to only the first engine.
|
|
This is useful when you prefer the reranking to not be affected by any potential overlap
|
|
of results from different engines.
|
|
|
|
.. code:: yaml
|
|
|
|
rerank:
|
|
remove_extra_engines: true
|
|
|
|
"""
|
|
|
|
id = "rerank"
|
|
default_on = False
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
self.stopword_langs = ['en', 'de', 'nl', 'fr', 'es', 'pt', 'it', 'ru', 'sv', 'no', 'zh']
|
|
self.remove_extra_engines = settings.get('rerank', {}).get('remove_extra_engines')
|
|
|
|
self.info = PluginInfo(
|
|
id=self.id,
|
|
name=gettext("Rerank plugin"),
|
|
description=gettext("""Rerank search results, ignoring original engine ranking"""),
|
|
preference_section="general",
|
|
is_allowed=bm25s is not None,
|
|
)
|
|
|
|
def post_search(self, request: "SXNG_Request", search: "SearchWithPlugins") -> EngineResults:
|
|
results = EngineResults()
|
|
|
|
if not bm25s:
|
|
return results
|
|
|
|
# pylint: disable=protected-access
|
|
results = search.result_container._merged_results
|
|
query = search.search_query.query
|
|
locale = search.search_query.locale
|
|
|
|
# Determine the stopwords based on the selected locale
|
|
stopwords = locale.language if locale and locale.language in self.stopword_langs else 'en'
|
|
|
|
retriever = bm25s.BM25()
|
|
result_tokens = bm25s.tokenize(
|
|
[
|
|
f"{result.get('title', '')} | {result.get('content', '')} | {result.get('url', '')}"
|
|
for result in results
|
|
],
|
|
stopwords=stopwords,
|
|
)
|
|
retriever.index(result_tokens)
|
|
|
|
query_tokens = bm25s.tokenize(query, stopwords=stopwords)
|
|
|
|
# Retrieve ranked indices of results based on the query tokens
|
|
indices = retriever.retrieve(query_tokens, k=len(results), return_as='documents', show_progress=False)
|
|
|
|
if self.remove_extra_engines:
|
|
# Only keep the main engine and set our ranking
|
|
for position, index in enumerate(indices[0]):
|
|
if 'positions' in results[index]:
|
|
results[index]['positions'] = [position + 1]
|
|
results[index]['engines'] = set([results[index]['engine']])
|
|
else:
|
|
# Overwrite all engine positions with the new ranking
|
|
# Results returned from multiple engines will still get a score boost
|
|
for position, index in enumerate(indices[0]):
|
|
if 'positions' in results[index]:
|
|
results[index]['positions'] = [position + 1] * len(results[index]['positions'])
|
|
|
|
return results
|