fix(retrieval): match the Qwen query prompt format
Drop the trailing space after "Query:" as in the model card and include the prompt in the query cache key. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
d60dd10c45
commit
f134c50d39
2 files changed
+8
-3
No files matched your search
@@ -133,6 +133,10 @@ class RetrievalServerTests(unittest.TestCase):
|
||||
self.assertNotEqual(MODULE.daemon_address(data_dir), old_address)
|
||||
self.assertNotEqual(MODULE.daemon_lock_path(data_dir), old_lock)
|
||||
|
||||
def test_query_prompt_matches_qwen_format(self):
|
||||
self.assertTrue(MODULE.QUERY_PROMPT.startswith('Instruct: '))
|
||||
self.assertTrue(MODULE.QUERY_PROMPT.endswith('\nQuery:'))
|
||||
|
||||
def test_preferred_device_uses_cuda_with_cpu_fallback(self):
|
||||
with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True))):
|
||||
self.assertEqual(MODULE.preferred_device(), 'cuda')
|
||||
|
||||
@@ -36,6 +36,8 @@ CLI_TIMEOUT_SECONDS = 300
|
||||
RRF_K = 60
|
||||
INDEX_VERSION = 3
|
||||
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
||||
# Qwen3-Embedding query format from the model card: no space after "Query:".
|
||||
QUERY_PROMPT = f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery:"
|
||||
TEXT_EXTENSIONS = {".c", ".cjs", ".cpp", ".cs", ".csproj", ".css", ".go", ".gradle", ".html", ".java", ".js", ".json", ".jsx", ".kt", ".kts",
|
||||
".md", ".mjs", ".props", ".ps1", ".psm1", ".py", ".razor", ".rs", ".scss", ".sh", ".sql", ".targets", ".toml", ".ts",
|
||||
".tsx", ".txt", ".xaml", ".xml", ".yaml", ".yml"}
|
||||
@@ -201,8 +203,7 @@ def shared_encoder():
|
||||
if _shared_encoder is None:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
_shared_encoder = SentenceTransformer(MODEL, revision=MODEL_REVISION, truncate_dim=EMBEDDING_DIMENSIONS,
|
||||
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
|
||||
device=preferred_device())
|
||||
prompts={"query": QUERY_PROMPT}, device=preferred_device())
|
||||
return _shared_encoder
|
||||
|
||||
|
||||
@@ -522,7 +523,7 @@ class Index:
|
||||
def cache_key(self, query: str, top_k: int):
|
||||
identity = json.dumps([
|
||||
query, top_k,
|
||||
RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION,
|
||||
RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION, QUERY_PROMPT,
|
||||
DENSE_CANDIDATES, LEXICAL_CANDIDATES, RERANK_CANDIDATES, FINAL_RESULTS, RRF_K,
|
||||
], ensure_ascii=False, separators=(",", ":"))
|
||||
return hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||
|
||||
Reference in new issue
Block a user