TBNilles/LEANN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: nits

Browse Source
This commit is contained in:
Andy Lee
2025-07-16 15:39:58 -07:00
parent 7b9406a3ea
commit 2a1a152073
4 changed files with 102 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
packages/leann-core/src/leann/api.py
View File

@@ -14,8 +14,7 @@ import torch
from .registry import BACKEND_REGISTRY
from .interface import LeannBackendFactoryInterface
# --- The Correct, Verified Embedding Logic from old_code.py ---
from .chat import get_llm
def compute_embeddings(
@@ -28,7 +27,7 @@ def compute_embeddings(
from sentence_transformers import SentenceTransformer
except ImportError as e:
raise RuntimeError(
f"sentence-transformers not available. Install with: pip install sentence-transformers"
"sentence-transformers not available. Install with: uv pip install sentence-transformers"
) from e
# Load model using sentence-transformers
@@ -61,7 +60,7 @@ def compute_embeddings_mlx(chunks: List[str], model_name: str) -> np.ndarray:
from mlx_lm.utils import load
except ImportError as e:
raise RuntimeError(
f"MLX or related libraries not available. Install with: uv pip install mlx mlx-lm"
"MLX or related libraries not available. Install with: uv pip install mlx mlx-lm"
) from e
print(
@@ -75,7 +74,7 @@ def compute_embeddings_mlx(chunks: List[str], model_name: str) -> np.ndarray:
all_embeddings = []
for chunk in chunks:
# Tokenize
token_ids = tokenizer.encode(chunk)
token_ids = tokenizer.encode(chunk) # type: ignore
# Convert to MLX array and add batch dimension
input_ids = mx.array([token_ids])
@@ -95,9 +94,6 @@ def compute_embeddings_mlx(chunks: List[str], model_name: str) -> np.ndarray:
return np.stack(all_embeddings)
# --- Core API Classes (Restored and Unchanged) ---
@dataclass
class SearchResult:
id: str
@@ -255,7 +251,7 @@ class LeannSearcher:
self.backend_impl = backend_factory.searcher(index_path, **final_kwargs)
def search(self, query: str, top_k: int = 5, **search_kwargs) -> List[SearchResult]:
print(f"🔍 DEBUG LeannSearcher.search() called:")
print("🔍 DEBUG LeannSearcher.search() called:")
print(f" Query: '{query}'")
print(f" Top_k: {top_k}")
print(f" Search kwargs: {search_kwargs}")
@@ -302,12 +298,13 @@ class LeannSearcher:
return enriched_results
from .chat import get_llm
class LeannChat:
def __init__(
self, index_path: str, llm_config: Optional[Dict[str, Any]] = None, enable_warmup: bool = False, **kwargs
self,
index_path: str,
llm_config: Optional[Dict[str, Any]] = None,
enable_warmup: bool = False,
**kwargs,
):
self.searcher = LeannSearcher(index_path, enable_warmup=enable_warmup, **kwargs)
self.llm = get_llm(llm_config)

7
packages/leann-core/src/leann/interface.py
View File

@@ -1,16 +1,17 @@
from abc import ABC, abstractmethod
import numpy as np
from typing import Dict, Any, Literal
from typing import Dict, Any, List, Literal
class LeannBackendBuilderInterface(ABC):
"""Backend interface for building indexes"""
@abstractmethod
def build(self, data: np.ndarray, index_path: str, **kwargs) -> None:
def build(self, data: np.ndarray, ids: List[str], index_path: str, **kwargs) -> None:
"""Build index
Args:
data: Vector data (N, D)
ids: List of string IDs for each vector
index_path: Path to save index
**kwargs: Backend-specific build parameters
"""
@@ -47,7 +48,7 @@ class LeannBackendSearcherInterface(ABC):
beam_width: Number of parallel search paths/IO requests per iteration
prune_ratio: Ratio of neighbors to prune via approximate distance (0.0-1.0)
recompute_embeddings: Whether to fetch fresh embeddings from server vs use stored PQ codes
pruning_strategy: PQ candidate selection strategy - "global", "local", or "proportional"
pruning_strategy: PQ candidate selection strategy - "global" (default), "local", or "proportional"
zmq_port: ZMQ port for embedding server communication
**kwargs: Backend-specific parameters