diff --git a/Dockerfile b/Dockerfile index 536746c..1767ada 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,7 +11,7 @@ COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt # Pre-download embedding model agar tidak diunduh ulang saat runtime -RUN python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('all-MiniLM-L6-v2')" +RUN python -c "from fastembed import TextEmbedding; TextEmbedding('sentence-transformers/all-MiniLM-L6-v2')" COPY src/ ./src/ diff --git a/requirements.txt b/requirements.txt index c83f8ec..52e2437 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,4 @@ python-dotenv==1.0.1 chromadb==0.5.3 pdfplumber==0.11.0 python-docx==1.1.2 -sentence-transformers==3.0.1 +fastembed==0.4.2 diff --git a/src/config.py b/src/config.py index 685641a..fc55307 100644 --- a/src/config.py +++ b/src/config.py @@ -11,4 +11,4 @@ ADMIN_TELEGRAM_ID = int(os.environ["ADMIN_TELEGRAM_ID"]) DB_PATH = os.getenv("DB_PATH", "/app/data/second_brain.db") CHROMA_PATH = os.getenv("CHROMA_PATH", "/app/data/chromadb") UPLOAD_PATH = os.getenv("UPLOAD_PATH", "/app/data/uploads") -EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "all-MiniLM-L6-v2") +EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2") diff --git a/src/rag.py b/src/rag.py index 91c185b..377fa09 100644 --- a/src/rag.py +++ b/src/rag.py @@ -1,10 +1,10 @@ import asyncio import chromadb -from sentence_transformers import SentenceTransformer +from fastembed import TextEmbedding from config import CHROMA_PATH, EMBEDDING_MODEL _chroma = chromadb.PersistentClient(path=CHROMA_PATH) -_embedding_model = SentenceTransformer(EMBEDDING_MODEL) +_embedding_model = TextEmbedding(model_name=EMBEDDING_MODEL) def _get_collection(name: str): @@ -12,8 +12,8 @@ def _get_collection(name: str): async def _embed(text: str) -> list: - embedding = await asyncio.to_thread(_embedding_model.encode, text) - return embedding.tolist() + result = await asyncio.to_thread(lambda: list(_embedding_model.embed([text]))[0]) + return result.tolist() async def add_memory(user_id: int, memory_id: int, content: str, scope: str = "private"):