Minor fixes, CI should build
This commit is contained in:
parent
519afe6709
commit
57f273e000
7 changed files with 30 additions and 6 deletions
2
.github/workflows/test.yml
vendored
2
.github/workflows/test.yml
vendored
|
|
@ -66,6 +66,8 @@ jobs:
|
||||||
key: huggingface-${{ runner.os }}-qwen-tokenizer-v1
|
key: huggingface-${{ runner.os }}-qwen-tokenizer-v1
|
||||||
- name: Pre-download tokenizer
|
- name: Pre-download tokenizer
|
||||||
run: uv run python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B')"
|
run: uv run python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B')"
|
||||||
|
- name: Pre-download cross-encoder test model
|
||||||
|
run: uv run python -c "from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')"
|
||||||
- name: Run tests with coverage
|
- name: Run tests with coverage
|
||||||
run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml
|
run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml
|
||||||
- name: Upload coverage to Codecov
|
- name: Upload coverage to Codecov
|
||||||
|
|
|
||||||
|
|
@ -58,6 +58,7 @@ async def download_models(
|
||||||
|
|
||||||
model_name = config.embeddings.model.name
|
model_name = config.embeddings.model.name
|
||||||
yield DownloadProgress(model=model_name, status="start")
|
yield DownloadProgress(model=model_name, status="start")
|
||||||
|
# Wrap in lambda: ty loses ParamSpec inference on third-party __init__.
|
||||||
await asyncio.to_thread(lambda: SentenceTransformer(model_name))
|
await asyncio.to_thread(lambda: SentenceTransformer(model_name))
|
||||||
yield DownloadProgress(model=model_name, status="done")
|
yield DownloadProgress(model=model_name, status="done")
|
||||||
except ImportError:
|
except ImportError:
|
||||||
|
|
|
||||||
|
|
@ -71,7 +71,12 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
try:
|
try:
|
||||||
from haiku.rag.reranking.cross_encoder import CrossEncoderReranker
|
from haiku.rag.reranking.cross_encoder import CrossEncoderReranker
|
||||||
|
|
||||||
return CrossEncoderReranker(config.reranking.model.name)
|
name = config.reranking.model.name
|
||||||
|
if not name:
|
||||||
|
raise ValueError(
|
||||||
|
"cross-encoder reranker requires name in reranking.model"
|
||||||
|
)
|
||||||
|
return CrossEncoderReranker(name)
|
||||||
except ImportError: # pragma: no cover
|
except ImportError: # pragma: no cover
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -33,8 +33,7 @@ class CrossEncoderReranker(RerankerBase):
|
||||||
return []
|
return []
|
||||||
|
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
loop = asyncio.get_running_loop()
|
rankings = await asyncio.to_thread(
|
||||||
rankings = await loop.run_in_executor(
|
lambda: self._reranker.rank(query, documents, top_k=top_n)
|
||||||
None, lambda: self._reranker.rank(query, documents, top_k=top_n)
|
|
||||||
)
|
)
|
||||||
return [(chunks[r["corpus_id"]], float(r["score"])) for r in rankings]
|
return [(chunks[r["corpus_id"]], float(r["score"])) for r in rankings]
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,5 @@
|
||||||
|
import asyncio
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from transformers import (
|
from transformers import (
|
||||||
AutoModel, # pyright: ignore[reportMissingImports]
|
AutoModel, # pyright: ignore[reportMissingImports]
|
||||||
|
|
@ -32,6 +34,8 @@ class JinaLocalReranker(RerankerBase): # pragma: no cover
|
||||||
|
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
results = self._reranker.rerank(query, documents, top_n=top_n)
|
results = await asyncio.to_thread(
|
||||||
|
lambda: self._reranker.rerank(query, documents, top_n=top_n)
|
||||||
|
)
|
||||||
|
|
||||||
return [(chunks[r["index"]], float(r["relevance_score"])) for r in results]
|
return [(chunks[r["index"]], float(r["relevance_score"])) for r in results]
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,5 @@
|
||||||
|
import asyncio
|
||||||
|
|
||||||
from mxbai_rerank import MxbaiRerankV2 # pyright: ignore[reportMissingImports]
|
from mxbai_rerank import MxbaiRerankV2 # pyright: ignore[reportMissingImports]
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
|
|
@ -22,7 +24,9 @@ class MxBAIReranker(RerankerBase):
|
||||||
|
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
results = self._client.rank(query=query, documents=documents, top_k=top_n)
|
results = await asyncio.to_thread(
|
||||||
|
lambda: self._client.rank(query=query, documents=documents, top_k=top_n)
|
||||||
|
)
|
||||||
reranked_chunks = []
|
reranked_chunks = []
|
||||||
for result in results:
|
for result in results:
|
||||||
original_chunk = chunks[result.index]
|
original_chunk = chunks[result.index]
|
||||||
|
|
|
||||||
|
|
@ -126,6 +126,15 @@ class TestGetReranker:
|
||||||
with pytest.raises(ValueError, match="vLLM reranker requires base_url"):
|
with pytest.raises(ValueError, match="vLLM reranker requires base_url"):
|
||||||
get_reranker(config)
|
get_reranker(config)
|
||||||
|
|
||||||
|
def test_cross_encoder_provider_without_name_raises_error(self):
|
||||||
|
config = AppConfig(
|
||||||
|
reranking=RerankingConfig(
|
||||||
|
model=ModelConfig(provider="cross-encoder", name="")
|
||||||
|
)
|
||||||
|
)
|
||||||
|
with pytest.raises(ValueError, match="cross-encoder reranker requires name"):
|
||||||
|
get_reranker(config)
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"provider, model_name, class_module, class_name, extra_model_kwargs, expected_attrs, env_vars",
|
"provider, model_name, class_module, class_name, extra_model_kwargs, expected_attrs, env_vars",
|
||||||
[
|
[
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue