44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
import os
|
|
|
|
from leann.api import LeannBuilder, LeannChat
|
|
|
|
# Define the path for our new MLX-based index
|
|
INDEX_PATH = "./mlx_diskann_index/leann"
|
|
|
|
if os.path.exists(INDEX_PATH + ".meta.json"):
|
|
print(f"Index already exists at {INDEX_PATH}. Skipping build.")
|
|
else:
|
|
print("Initializing LeannBuilder with MLX support...")
|
|
# 1. Configure LeannBuilder to use MLX
|
|
builder = LeannBuilder(
|
|
backend_name="hnsw",
|
|
embedding_model="mlx-community/Qwen3-Embedding-0.6B-4bit-DWQ",
|
|
embedding_mode="mlx",
|
|
)
|
|
|
|
# 2. Add documents
|
|
print("Adding documents...")
|
|
docs = [
|
|
"MLX is an array framework for machine learning on Apple silicon.",
|
|
"It was designed by Apple's machine learning research team.",
|
|
"The mlx-community organization provides pre-trained models in MLX format.",
|
|
"It supports operations on multi-dimensional arrays.",
|
|
"Leann can now use MLX for its embedding models.",
|
|
]
|
|
for doc in docs:
|
|
builder.add_text(doc)
|
|
|
|
# 3. Build the index
|
|
print(f"Building the MLX-based index at: {INDEX_PATH}")
|
|
builder.build_index(INDEX_PATH)
|
|
print("\nSuccessfully built the index with MLX embeddings!")
|
|
print(f"Check the metadata file: {INDEX_PATH}.meta.json")
|
|
|
|
|
|
chat = LeannChat(index_path=INDEX_PATH)
|
|
# add query
|
|
query = "MLX is an array framework for machine learning on Apple silicon."
|
|
print(f"Query: {query}")
|
|
response = chat.ask(query, top_k=3, recompute_beighbor_embeddings=True, complexity=3, beam_width=1)
|
|
print(f"Response: {response}")
|