factor out embeddings.py

Author: jordanstephens

app/Dockerfile

@@ -36,8 +36,10 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Set the environment variable for the sentence transformers model
 ENV SENTENCE_TRANSFORMERS_HOME="/root/.cache/sentence_transformers"
 
+COPY ./embeddings.py /app/embeddings.py
+
 # Preload the sentence transformer model to cache
-RUN python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('all-MiniLM-L6-v2')"
+RUN python embeddings.py
 
 # Copy the application source code into the container
 COPY . /app

app/embeddings.py

@@ -0,0 +1,15 @@
+import os
+from sentence_transformers import SentenceTransformer
+
+def load_model():
+    """
+    Load the SentenceTransformer model from the specified path.
+    The model path is determined by the SENTENCE_TRANSFORMERS_HOME environment variable.
+    """
+    # model_path = os.getenv("SENTENCE_TRANSFORMERS_HOME", "./models/sentence-transformers")
+    # model = SentenceTransformer(f"{model_path}/models--sentence-transformers--all-MiniLM-L6-v2")
+    model = SentenceTransformer("all-MiniLM-L6-v2")
+    return model
+
+if __name__ == "__main__":
+    load_model()

app/rag_system.py

@@ -6,6 +6,7 @@
 from sentence_transformers import SentenceTransformer
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
+from embeddings import load_model
 import traceback
 
 openai.api_base = os.getenv("OPENAI_BASE_URL")
@@ -15,7 +16,7 @@ class RAGSystem:
     def __init__(self, knowledge_base_path='./data/knowledge_base.json'):
         self.knowledge_base_path = knowledge_base_path
         self.knowledge_base = self.load_knowledge_base()
-        self.model = SentenceTransformer('all-MiniLM-L6-v2')
+        self.model = load_model()
         self.doc_embeddings = self.embed_knowledge_base()
         self.conversation_history = []