Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

HybridSearchEngine

Where I implement my idea about finding and searching information.

Python Pydantic Status Models

This is my personal lab for one question: how do you actually find the right piece of information inside a large, messy corpus?

Every retrieval idea I read about, argue with, or come up with ends up here — implemented for real, wired into a runnable pipeline, and tested.


The idea

Search is not a single ranking function. It is a chain of decisions:

  1. Understand the question. One user question is often several real questions.
  2. Look in more than one place. Dense vectors, summary vectors, BM25, entities, relations, document numbers, structural location — each channel sees something the others miss.
  3. Fuse without lying. A cosine score and a BM25 score are not comparable numbers. Fuse by rank, not by raw score.
  4. Re-read carefully. A cross-encoder gets the final word on what answers the sub-question.
  5. Return evidence, not blobs. An answer that cannot point to Điều 20 of document X, page 15 is not an answer.

Every stage sits behind a Protocol in interfaces.py, so any part can be replaced and measured.

What is implemented

Idea Where it lives
Query classification, rewrite, decomposition into sub-queries modules/decompose/
Multi-channel retrieval (dense content + dense summary + sparse) modules/retrieval/
Per-sub-query Reciprocal Rank Fusion core/fusion.py
Cross-encoder rerank with per-query threshold and top-K modules/rerank/
Evidence assembly, union/dedup, optional base score ordering pipelines/evidence.py
Hierarchical search carrying seeds between layers pipelines/deep.py
Embedding and rerank inference: HuggingFace, ONNX, Triton ai_inferences/
Pluggable vector and text stores connectors/
Typed, immutable, validated domain models models.py

Backends: Elasticsearch, FAISS, Milvus, Qdrant.

How it flows

Q ──> decompose into q_i
      └─> for each q_i: dense(content) + dense(summary) + sparse(BM25)
          └─> RRF per q_i
              └─> cross-encoder(q_i, chunk)
                  └─> threshold + top-K per q_i
                      └─> Evidence ──> union / dedup chunks
                                       └─> optional base score for final ordering

Layout

src/hybrid_search/
├── ai_inferences/  # Embedding and rerank engines: HuggingFace, ONNX, Triton
│   ├── embedding.py
│   └── reranking.py
├── connectors/     # DatabaseConnector: Elasticsearch, FAISS, Milvus, Qdrant
├── core/           # Dedupe, merge, fusion, score normalization — pure Python
├── modules/
│   ├── decompose/  # Query classification, rewrite, decomposition
│   ├── retrieval/  # Dense / sparse / metadata retrieval + fusion
│   └── rerank/     # Per-query threshold, top-K, optional base score
├── pipelines/      # decompose -> retrieval -> rerank
├── models.py       # Every Pydantic model: domain objects and configs, validated, immutable
└── interfaces.py   # Protocols for LLM / embed / retrieve / rerank providers

Inference engines

The pipeline never imports torch, onnxruntime or a Triton client. It calls two methods:

await embedding_engine.embed(text, profile="query")   # -> vector
await rerank_engine.score(query, chunks)              # -> one score in [0, 1] per chunk

Every engine in ai_inferences/ exposes exactly those, so changing runtime or checkpoint is a constructor change and nothing above it moves:

Runtime Embedding Rerank
HuggingFace / PEFT HFEmbed HFRerank
ONNX Runtime OnnxEmbed OnnxRerank
Triton Inference Server (AI-Serving) TritonEmbed TritonRerank

BaseEmbedding owns profile routing, batching and L2 normalization; BaseReranking owns the prompt format, batching and squashing raw logits into [0, 1] — a new runtime only implements infer_batch.

How the layer fits together, what profile selects, and the pooling traps are written up in src/hybrid_search/ai_inferences/README.md.

from hybrid_search.ai_inferences import HFRerank, OnnxRerank, TritonRerank

# local, fine-tuned LoRA
scorer = HFRerank("hoailebads/Qwen3-Reranker-0.6B-VLSP-Legal-LoRA")

# exported graph — same call
scorer = OnnxRerank("reranker.onnx", tokenizer="Qwen/Qwen3-Reranker-0.6B")

# served — same call
scorer = TritonRerank("localhost:8001", backend="tensorrt")
pip install -e '.[hf]'       # torch + transformers + peft
pip install -e '.[onnx]'     # onnxruntime
pip install -e '.[triton]'   # tritonclient[grpc]

Serving with AI-Serving

TritonEmbed and TritonRerank speak that stack's contract out of the box: gRPC on localhost:8001, one text input as a 1-D BYTES tensor, the output requested explicitly, and model names built from the backend.

Embedding Rerank
Model name embedding_onnx / embedding_tensorrt reranker_onnx / reranker_tensorrt
Input text (BYTES, [-1]) text (BYTES, [-1])
Output embeddings (FP32, [-1, -1]) scores (FP32, [-1])
from hybrid_search.ai_inferences import TritonEmbed, TritonRerank

embedder = TritonEmbed("localhost:8001", backend="tensorrt")
scorer = TritonRerank("localhost:8001", backend="tensorrt")

The ensemble does tokenizer → core → pooling itself, and its pooling backend already mean-pools and L2-normalizes, so TritonEmbed defaults to normalize=False instead of normalizing a second time. Two ensembles for the dual-adapter model? Pass passage_model_name= and the passage profile routes there.

OnnxEmbed covers the other half of that repo: converter/convert_onnx.py exports token_embeddings (last hidden state), so pooling happens client side — it defaults to pooling="mean" to match the serving pooling backend. HFEmbed defaults to pooling="last" because that is what the Qwen3 checkpoints were trained with; keep the two consistent or the exported graph will not reproduce local scores.

My models

The fine-tuned checkpoints live at huggingface.co/hoailebads, targeting Vietnamese legal retrieval (VLSP / Zalo Legal):

Model Task Base Form
Qwen3-Reranker-0.6B-VLSP-Legal-LoRA Text ranking Qwen/Qwen3-Reranker-0.6B LoRA r=32, alpha=64, SEQ_CLS
Qwen3-Reranker-8B-VLSP-Legal-LoRA Text ranking Qwen/Qwen3-Reranker-8B LoRA adapter
Qwen3-Embedding-0.6B-VLSP-Legal-Retrieval-LoRA Feature extraction Qwen/Qwen3-Embedding-0.6B Dual LoRA: query_adapter + passage_adapter, dim 1024
Qwen3-Embedding-0.6B-Zalo-Legal-Retrieval-LoRA Feature extraction Qwen/Qwen3-Embedding-0.6B LoRA adapter

HFEmbed and HFRerank load a plain checkpoint, a single LoRA adapter, or the dual-adapter layout — it inspects the repo and picks the right one:

from hybrid_search.ai_inferences import HFEmbed

# stock model, mean pooling
embedder = HFEmbed("intfloat/multilingual-e5-base", pooling="mean")

# my dual-LoRA model: query_adapter / passage_adapter switched per call
embedder = HFEmbed("hoailebads/Qwen3-Embedding-0.6B-VLSP-Legal-Retrieval-LoRA")

await embedder.embed(question, profile="query")     # query_adapter
await embedder.embed(article, profile="passage")    # passage_adapter

Encoding a query with the passage adapter silently wrecks recall, so profile is what selects the adapter. Build the retrieval module with embedding_profile="query" and index passages with profile="passage".

The rerankers score with a sequence-classification head (num_labels=1, scalar logit at the last token), not the yes/no generative trick of the stock Qwen3 reranker. Prompt:

Instruct: {instruction}
query: {query}
document: {chunk text}<eos>

On the VLSP eval set (219 queries, 67,561 chunks from 59,628 articles), reranking with the 0.6B adapter gives R@1 55.18 / R@3 71.39 / R@5 78.23 / R@10 85.12+9.82 R@1 and +3.99 R@10 over retrieval alone. Logits are comparable within one query only, which is exactly how RerankModule uses them: per sub-query threshold and top-K, then fusion by rank.

Wiring a pipeline

pip install -e '.[hf,faiss,elasticsearch]'
PYTHONPATH=src python example/evidence_search.py
from hybrid_search.ai_inferences import HFEmbed, HFRerank
from hybrid_search.connectors import (
    ElasticsearchConfig,
    ElasticsearchDatabaseConnector,
    FaissDatabaseConnector,
)
from hybrid_search.modules import DecomposeModule, RerankModule, RetrievalModule
from hybrid_search.pipelines import EvidenceSearch

embedder = HFEmbed("hoailebads/Qwen3-Embedding-0.6B-VLSP-Legal-Retrieval-LoRA")
dense = FaissDatabaseConnector(chunks)
sparse = ElasticsearchDatabaseConnector(
    ElasticsearchConfig(hosts=("http://localhost:9200",), index_name="chunks")
)
await dense.connect()
await sparse.connect()

searcher = EvidenceSearch(
    decompose_module=DecomposeModule(),
    retrieval_module=RetrievalModule(
        embedder=embedder,
        dense_database=dense,
        sparse_database=sparse,
        chunks=chunks,
        embedding_profile="query",
    ),
    rerank_module=RerankModule(
        scorer=HFRerank("hoailebads/Qwen3-Reranker-0.6B-VLSP-Legal-LoRA")
    ),
)
response = await searcher.search("thời hạn hợp đồng lao động")

FaissDatabaseConnector keeps separate content and summary indexes, cosine similarity, metadata post-filter and index factories (Flat, HNSW,Flat, IVF...). For big corpora, shard by domain or document type instead of scanning the whole index behind a filter.

ElasticsearchDatabaseConnector runs multi_match over text, summary, entities and document numbers. Raw BM25 stays in provenance only — fusion works on rank, so BM25 is never added to a cosine score.

The contracts

Everything is injected through interfaces.py:

Contract Implemented by
QueryAnalyzer An OpenAI-compatible model with structured output
Embedder HFEmbed, OnnxEmbed, TritonEmbed
ChunkScorer HFRerank, OnnxRerank, TritonRerank
DatabaseConnector Elasticsearch / FAISS / Milvus / Qdrant
FullTextRetriever Elasticsearch

Design rules I hold myself to

  • No mutable default arguments, ever.
  • Never mutate what a provider returned; copy, then transform.
  • Never pair a query with the wrong embedding — or the wrong adapter.
  • No hard-coded endpoints or credentials in pipeline code; they live in connector config.
  • No fake backends shipped in the library. Test doubles belong in tests/.
  • Every data/config object is a Pydantic model: validated, frozen, no extra fields — so an API payload goes through model_validate(...) and results come back via model_dump() with types intact end to end.

Notes and deeper write-ups

Where this is going

  • Calibrated, comparable scores across channels
  • Answerability: knowing when the corpus does not contain the answer
  • Constraint matching (time, jurisdiction, entity) as a first-class ranking signal
  • Multi-aspect EvidenceScore instead of a single number
  • Benchmarks on real corpora, so ideas get judged by numbers and not by taste

Tests

PYTHONPATH=src:tests python -m unittest discover -s tests -v

This repository is a workbench, not a product. If something in it is useful to you, take it.

About

Where i implement my idea about finding and searching information

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages