wehub-resource-sync
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
"""
|
|
Sample code to define a query engine with post processors and save it with model-from-code logging.
|
|
|
|
Ref: https://qdrant.tech/documentation/quickstart/
|
|
"""
|
|
|
|
from llama_index.core import Document, QueryBundle, VectorStoreIndex
|
|
from llama_index.core.postprocessor import LLMRerank
|
|
from llama_index.core.postprocessor.types import BaseNodePostprocessor
|
|
from llama_index.core.schema import NodeWithScore
|
|
|
|
import mlflow
|
|
|
|
index = VectorStoreIndex.from_documents(documents=[Document.example()])
|
|
|
|
# Postprocessor 1: Reranker
|
|
reranker = LLMRerank(
|
|
choice_batch_size=5,
|
|
top_n=2,
|
|
)
|
|
|
|
|
|
# Postprocessor 2: custom postprocessor
|
|
class CustomNodePostprocessor(BaseNodePostprocessor):
|
|
call_count: int = 0
|
|
|
|
def _postprocess_nodes(
|
|
self, nodes: list[NodeWithScore], query_bundle: QueryBundle | None
|
|
) -> list[NodeWithScore]:
|
|
# subtracts 1 from the score
|
|
self.call_count += 1
|
|
return nodes
|
|
|
|
|
|
query_engine = index.as_query_engine(
|
|
similarity_top_k=5,
|
|
node_postprocessors=[reranker, CustomNodePostprocessor()],
|
|
)
|
|
|
|
|
|
mlflow.models.set_model(query_engine)
|