Import dense-retrieval

Makefile

@@ -8,6 +8,7 @@ TEMPLATES_DIR:=templates
 .PHONY: lint
 lint:
 	python -m ruff check --fix --unsafe-fixes --show-fixes
+	python -m mypy src/dense-retrieval/src --explicit-package-bases --namespace-packages
 	python -m mypy src/amazon-product-search/src --explicit-package-bases --namespace-packages
 	python -m mypy src/training/src --explicit-package-bases --namespace-packages
 

pyproject.toml

@@ -38,7 +38,8 @@ extend-select = ["I"]
 [tool.ruff.lint.isort]
 known-first-party = [
     "amazon_product_search",
-    "amazon_product_search_dense_retrieval",
+    "dense_retrieval",
+    "training",
 ]
 
 [tool.ruff.lint.pydocstyle]

src/amazon-product-search/pyproject.toml

@@ -26,7 +26,7 @@ unidic-lite = "^1.0.8"
 polars = {version = "^0.16.0", extras = ["numpy", "pandas", "pyarrow"]}
 types-requests = "^2.28.11.17"
 docker = "^7.0.0"
-amazon-product-search-dense-retrieval = {git = "https://github.com/rejasupotaro/amazon-product-search-dense-retrieval.git"}
+dense-retrieval = { path = "../dense-retrieval", develop = true }
 gcsfs = "^2023.6.0"
 jinja2 = "^3.1.2"
 google-cloud-bigquery = {extras = ["pandas"], version = "^3.12.0"}

src/amazon-product-search/src/amazon_product_search/core/es/query_builder.py

@@ -8,7 +8,7 @@
 from amazon_product_search.core.retrieval.query_vector_cache import QueryVectorCache
 from amazon_product_search.core.source import Locale
 from amazon_product_search.core.synonyms.synonym_dict import SynonymDict, expand_synonyms
-from amazon_product_search_dense_retrieval.encoders import SBERTEncoder
+from dense_retrieval.encoders import SBERTEncoder
 
 
 class QueryBuilder:

src/amazon-product-search/src/amazon_product_search/core/synonyms/filters/similarity_filter.py

@@ -5,7 +5,7 @@
 from tqdm import tqdm
 
 from amazon_product_search.constants import HF
-from amazon_product_search_dense_retrieval.encoders import SBERTEncoder
+from dense_retrieval.encoders import SBERTEncoder
 
 
 class SimilarityFilter:

src/amazon-product-search/src/amazon_product_search/core/vespa/query_builder.py

@@ -6,7 +6,7 @@
 from amazon_product_search.core.retrieval.query_vector_cache import QueryVectorCache
 from amazon_product_search.core.source import Locale
 from amazon_product_search.core.synonyms.synonym_dict import SynonymDict
-from amazon_product_search_dense_retrieval.encoders import SBERTEncoder
+from dense_retrieval.encoders import SBERTEncoder
 
 Operator = Literal["and", "weakAnd"]
 

src/amazon-product-search/src/amazon_product_search/indexing/query_pipeline.py

@@ -16,8 +16,8 @@
 from amazon_product_search.core.nlp.normalizer import normalize_query
 from amazon_product_search.core.source import Locale
 from amazon_product_search.indexing.options import IndexerOptions
-from amazon_product_search_dense_retrieval.encoders import SBERTEncoder
-from amazon_product_search_dense_retrieval.encoders.modules.pooler import PoolingMode
+from dense_retrieval.encoders import SBERTEncoder
+from dense_retrieval.encoders.modules.pooler import PoolingMode
 
 
 def get_input_source(data_dir: str, locale: Locale, nrows: int = -1) -> PTransform:

src/amazon-product-search/src/amazon_product_search/indexing/transforms/encode_product.py

@@ -11,8 +11,8 @@
     InferInput,
 )
 
-from amazon_product_search_dense_retrieval.encoders import SBERTEncoder
-from amazon_product_search_dense_retrieval.encoders.modules.pooler import PoolingMode
+from dense_retrieval.encoders import SBERTEncoder
+from dense_retrieval.encoders.modules.pooler import PoolingMode
 
 
 def _product_to_text(product: Dict[str, Any], fields: list[str]) -> str:

src/amazon-product-search/tasks/vespa_tasks.py

@@ -5,7 +5,7 @@
 import amazon_product_search.core.vespa.service as vespa_service
 from amazon_product_search.constants import HF, VESPA_DIR
 from amazon_product_search.core.vespa.vespa_client import VespaClient
-from amazon_product_search_dense_retrieval.encoders import SBERTEncoder
+from dense_retrieval.encoders import SBERTEncoder
 
 """
 To run Vespa locally, execute the following commands:

src/dense-retrieval/README.md

@@ -0,0 +1,19 @@
+# Dense Retrieval - Amazon Product Search
+
+## Installation
+
+```shell
+$ pyenv install 3.11.8
+$ pyenv local 3.11.8
+$ pip install poetry
+$ poetry env use python
+$ poetry install
+```
+
+The following libraries are necessary for Japanese text processing.
+
+```shell
+# For macOS
+$ brew install mecab mecab-ipadic
+$ poetry run python -m unidic download
+```