feat(qwen): enforce 1024-dimensional embeddings
This commit is contained in:
1 parent
459c520201
commit
7be729cb22
2 files changed
+11
-5
No files matched your search
@@ -16,6 +16,10 @@ MODULE = importlib.util.module_from_spec(SPEC)
|
|||||||
SPEC.loader.exec_module(MODULE)
|
SPEC.loader.exec_module(MODULE)
|
||||||
|
|
||||||
|
|
||||||
|
def embedding():
|
||||||
|
return array('f', [1.0] + [0.0] * (MODULE.EMBEDDING_DIMENSIONS - 1))
|
||||||
|
|
||||||
|
|
||||||
class RetrievalServerTests(unittest.TestCase):
|
class RetrievalServerTests(unittest.TestCase):
|
||||||
@unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.')
|
@unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.')
|
||||||
def test_real_models(self):
|
def test_real_models(self):
|
||||||
@@ -52,7 +56,7 @@ class RetrievalServerTests(unittest.TestCase):
|
|||||||
def encode(values, **kwargs):
|
def encode(values, **kwargs):
|
||||||
self.assertEqual(kwargs.get('prompt_name'), 'query' if values == ['query'] else None)
|
self.assertEqual(kwargs.get('prompt_name'), 'query' if values == ['query'] else None)
|
||||||
encoded.extend(values)
|
encoded.extend(values)
|
||||||
return [array('f', [1.0]) for _ in values]
|
return [embedding() for _ in values]
|
||||||
|
|
||||||
for index in indexes:
|
for index in indexes:
|
||||||
index.encoder = lambda: SimpleNamespace(encode=encode)
|
index.encoder = lambda: SimpleNamespace(encode=encode)
|
||||||
@@ -88,7 +92,7 @@ class RetrievalServerTests(unittest.TestCase):
|
|||||||
root = Path(temporary)
|
root = Path(temporary)
|
||||||
(root / 'a.md').write_text('before', encoding='utf-8')
|
(root / 'a.md').write_text('before', encoding='utf-8')
|
||||||
index = MODULE.Index(root, root / 'data')
|
index = MODULE.Index(root, root / 'data')
|
||||||
index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [array('f', [1.0]) for _ in values])
|
index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [embedding() for _ in values])
|
||||||
try:
|
try:
|
||||||
index.rebuild()
|
index.rebuild()
|
||||||
(root / 'a.md').write_text('after', encoding='utf-8')
|
(root / 'a.md').write_text('after', encoding='utf-8')
|
||||||
@@ -97,7 +101,7 @@ class RetrievalServerTests(unittest.TestCase):
|
|||||||
def encode(values, **kwargs):
|
def encode(values, **kwargs):
|
||||||
if values == ['fail']:
|
if values == ['fail']:
|
||||||
raise RuntimeError('encoder failed')
|
raise RuntimeError('encoder failed')
|
||||||
return [array('f', [1.0]) for _ in values]
|
return [embedding() for _ in values]
|
||||||
|
|
||||||
index.encoder = lambda: SimpleNamespace(encode=encode)
|
index.encoder = lambda: SimpleNamespace(encode=encode)
|
||||||
with self.assertRaisesRegex(RuntimeError, 'encoder failed'):
|
with self.assertRaisesRegex(RuntimeError, 'encoder failed'):
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ import threading
|
|||||||
|
|
||||||
MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
||||||
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
|
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
|
||||||
|
EMBEDDING_DIMENSIONS = 1024
|
||||||
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
||||||
TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"}
|
TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"}
|
||||||
|
|
||||||
@@ -31,7 +32,7 @@ class Index:
|
|||||||
self.root = root.resolve()
|
self.root = root.resolve()
|
||||||
self.data = data.resolve()
|
self.data = data.resolve()
|
||||||
self.data.mkdir(parents=True, exist_ok=True)
|
self.data.mkdir(parents=True, exist_ok=True)
|
||||||
identity = os.path.normcase(str(self.root)) + "\n" + MODEL
|
identity = os.path.normcase(str(self.root)) + "\n" + MODEL + "\n" + str(EMBEDDING_DIMENSIONS)
|
||||||
database = hashlib.sha256(identity.encode()).hexdigest() + ".sqlite3"
|
database = hashlib.sha256(identity.encode()).hexdigest() + ".sqlite3"
|
||||||
self.db = sqlite3.connect(self.data / database, check_same_thread=False, timeout=30)
|
self.db = sqlite3.connect(self.data / database, check_same_thread=False, timeout=30)
|
||||||
self.db.execute("create table if not exists chunks (path text, text text, vector blob)")
|
self.db.execute("create table if not exists chunks (path text, text text, vector blob)")
|
||||||
@@ -62,7 +63,8 @@ class Index:
|
|||||||
def encoder(self):
|
def encoder(self):
|
||||||
if self.model is None:
|
if self.model is None:
|
||||||
from sentence_transformers import SentenceTransformer
|
from sentence_transformers import SentenceTransformer
|
||||||
self.model = SentenceTransformer(MODEL, prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "})
|
self.model = SentenceTransformer(MODEL, truncate_dim=EMBEDDING_DIMENSIONS,
|
||||||
|
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "})
|
||||||
return self.model
|
return self.model
|
||||||
|
|
||||||
def reranker(self):
|
def reranker(self):
|
||||||
|
|||||||
Reference in new issue
Block a user