made local rag needs optional (prevents having to have torch / sentence-transformers if you never use local embedding)
This commit is contained in:
+9
-14
@@ -16,6 +16,7 @@ from src.file_cache import ASTParser
|
||||
_SENTENCE_TRANSFORMERS = None
|
||||
_GOOGLE_GENAI = None
|
||||
_CHROMADB = None
|
||||
LOCAL_RAG_INSTALL_HINT = "Local RAG embeddings require sentence-transformers. Install with manual_slop[local-rag] to use local embeddings."
|
||||
|
||||
|
||||
def _get_sentence_transformers():
|
||||
@@ -24,6 +25,10 @@ def _get_sentence_transformers():
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
_SENTENCE_TRANSFORMERS = SentenceTransformer
|
||||
except ModuleNotFoundError as e:
|
||||
if e.name == "sentence_transformers":
|
||||
raise ImportError(LOCAL_RAG_INSTALL_HINT) from e
|
||||
raise
|
||||
except Exception as e:
|
||||
sys.stderr.write(f"FAILED to import sentence_transformers: {e}\n")
|
||||
sys.stderr.flush()
|
||||
@@ -52,22 +57,12 @@ class BaseEmbeddingProvider:
|
||||
|
||||
class LocalEmbeddingProvider(BaseEmbeddingProvider):
|
||||
def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
|
||||
self.model = None
|
||||
try:
|
||||
ST = _get_sentence_transformers()
|
||||
if ST:
|
||||
self.model = ST(model_name)
|
||||
except Exception as e:
|
||||
sys.stderr.write(f"LocalEmbeddingProvider failed to load model {model_name}: {e}. Using dummy embeddings.\n")
|
||||
sys.stderr.flush()
|
||||
ST = _get_sentence_transformers()
|
||||
self.model = ST(model_name)
|
||||
|
||||
def embed(self, texts: List[str]) -> List[List[float]]:
|
||||
if self.model:
|
||||
embeddings = self.model.encode(texts)
|
||||
return embeddings.tolist()
|
||||
else:
|
||||
# Dummy embeddings (384 dims for all-MiniLM-L6-v2)
|
||||
return [[0.0] * 384 for _ in texts]
|
||||
embeddings = self.model.encode(texts)
|
||||
return embeddings.tolist()
|
||||
|
||||
class GeminiEmbeddingProvider(BaseEmbeddingProvider):
|
||||
def __init__(self, model_name: str = 'gemini-embedding-001'):
|
||||
|
||||
Reference in New Issue
Block a user