Julian Carlson
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
"""Reranking via Qwen3-Reranker over Ollama."""
|
|
import requests
|
|
import os
|
|
import re
|
|
from .config import CONFIG
|
|
|
|
RERANK_MODEL = os.getenv("KB_RERANK_MODEL", "dengcao/Qwen3-Reranker-4B:Q5_K_M")
|
|
|
|
|
|
def score_relevance(query: str, document: str) -> float:
|
|
"""Get a relevance score in [0, 1] for (query, document) pair."""
|
|
if len(document) > 4000:
|
|
document = document[:4000] + "..."
|
|
|
|
# Pre-fill <think></think> as empty so the model jumps straight to answer
|
|
prompt = (
|
|
f"<|im_start|>system\n"
|
|
f"You evaluate whether a document is relevant to a query about quantitative trading. "
|
|
f"Be strict: the document must specifically address the query. "
|
|
f"Output only \"yes\" or \"no\".<|im_end|>\n"
|
|
f"<|im_start|>user\n"
|
|
f"<Query>: {query}\n"
|
|
f"<Document>: {document}<|im_end|>\n"
|
|
f"<|im_start|>assistant\n"
|
|
f"<think>\n\n</think>\n\n"
|
|
)
|
|
|
|
try:
|
|
resp = requests.post(
|
|
f"{CONFIG.ollama_url}/api/generate",
|
|
json={
|
|
"model": RERANK_MODEL,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"options": {
|
|
"num_predict": 5,
|
|
"temperature": 0.0,
|
|
"top_p": 1.0,
|
|
},
|
|
"raw": True,
|
|
},
|
|
timeout=60,
|
|
)
|
|
resp.raise_for_status()
|
|
text = resp.json().get("response", "").strip().lower()
|
|
text = re.sub(r"</?think>", "", text).strip()
|
|
resp.raise_for_status()
|
|
text = resp.json().get("response", "").strip()
|
|
text = re.sub(r"</?think>", "", text).strip()
|
|
|
|
# DEBUG: print every response so we can see what's happening
|
|
print(f"DEBUG rerank response: {text[:80]!r}", flush=True)
|
|
|
|
text = text.lower()
|
|
|
|
if text.startswith("yes"):
|
|
return 1.0
|
|
elif text.startswith("no"):
|
|
return 0.0
|
|
else:
|
|
print(f"Rerank ambiguous response: {text[:100]!r}", flush=True)
|
|
return 0.5
|
|
except Exception as e:
|
|
print(f"Rerank error: {e}", flush=True)
|
|
return 0.5
|
|
|
|
|
|
def rerank(query: str, candidates: list, top_k: int = 10) -> list:
|
|
if not candidates:
|
|
return []
|
|
|
|
scored = []
|
|
for cand in candidates:
|
|
chunk_meta = cand.get("chunk_metadata") or {}
|
|
doc_text = chunk_meta.get("raw_chunk") or cand.get("content", "")
|
|
score = score_relevance(query, doc_text)
|
|
cand_with_score = dict(cand)
|
|
cand_with_score["rerank_score"] = score
|
|
cand_with_score["original_similarity"] = cand.get("similarity")
|
|
scored.append(cand_with_score)
|
|
|
|
scored.sort(key=lambda x: x["rerank_score"], reverse=True)
|
|
return scored[:top_k]
|