Update src/apps/utils/embeddings.py

src/apps/utils/embeddings.py

@@ -1,22 +1,27 @@
-import os
-from transformers import AutoModel
-from numpy.linalg import norm
-
-cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
-_model = None
-
-def get_model():
-    global _model
-    if _model is None:
-        BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
-        model_path = os.path.join(BASE_DIR, "volumes", "models", "jina-embeddings-v2-base-en")
-        print(f"Loading the model weights from: {model_path}")
-        _model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
-    return _model
-
-def get_embeddings(text:list):
-    model = get_model()
-    embeddings = model.encode(text)
-    normalized_embeddings = embeddings/norm(embeddings[0])
-    return normalized_embeddings
-
+import os
+from transformers import AutoModel
+from numpy.linalg import norm
+
+cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
+_model = None
+
+def get_model():
+    global _model
+    if _model is None:
+        BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+        model_path = os.path.join(BASE_DIR, "volumes", "models", "jina-embeddings-v2-base-en")
+        
+        if os.path.exists(model_path):
+            print(f"Loading the model weights from local path: {model_path}")
+            _model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
+        else:
+            print("Local model weights not found. Downloading from Hugging Face Hub (jinaai/jina-embeddings-v2-base-en)...")
+            _model = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-en", trust_remote_code=True)
+    return _model
+
+def get_embeddings(text:list):
+    model = get_model()
+    embeddings = model.encode(text)
+    normalized_embeddings = embeddings/norm(embeddings[0])
+    return normalized_embeddings
+