Compute embedding using local resource directly

src/utilities/embedding.py

@@ -4,21 +4,14 @@ from typing import Any, List
 
 import torch
 from langchain_core.embeddings import Embeddings
-from langchain_huggingface import (
-    HuggingFaceEmbeddings,
-    HuggingFaceEndpointEmbeddings,
-)
+from langchain_huggingface import HuggingFaceEmbeddings
 from pydantic import BaseModel, Field
 
-
 class CustomEmbedding(BaseModel, Embeddings):
     """
-    Custom embedding class that supports both hosted and CPU embeddings.
+    Custom embedding class that supports CPU embeddings.
     """
 
-    hosted_embedding: HuggingFaceEndpointEmbeddings = Field(
-        default_factory=lambda: None
-    )
     cpu_embedding: HuggingFaceEmbeddings = Field(default_factory=lambda: None)
     matryoshka_dim: int = Field(default=256)
 
@@ -71,23 +64,8 @@ class CustomEmbedding(BaseModel, Embeddings):
         super().__init__(**kwargs)
         query_instruction = self.get_instruction()
         self.matryoshka_dim = matryoshka_dim
-        if torch.cuda.is_available():
-            logging.info("CUDA is available")
-            self.hosted_embedding = self.get_hf_embedd()
-            self.cpu_embedding = self.hosted_embedding
-        else:
-            logging.info("CUDA is not available")
-            self.hosted_embedding = HuggingFaceEndpointEmbeddings(
-                model=os.getenv("HF_MODEL"),
-                model_kwargs={
-                    "encode_kwargs": {
-                        "normalize_embeddings": True,
-                        "prompt": query_instruction,
-                    }
-                },
-                huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
-            )
-            self.cpu_embedding = self.get_hf_embedd()
+        logging.info("Initializing CPU embedding")
+        self.cpu_embedding = self.get_hf_embedd()
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
         """
@@ -100,10 +78,10 @@ class CustomEmbedding(BaseModel, Embeddings):
             List[List[float]]: List of embedded document vectors.
         """
         try:
-            embed = self.hosted_embedding.embed_documents(texts)
-        except Exception as e:
-            logging.warning(f"Issue with batch hosted embedding, moving to CPU: {e}")
             embed = self.cpu_embedding.embed_documents(texts)
+        except Exception as e:
+            logging.warning(f"Issue with CPU embedding: {e}")
+            embed = []
         return (
             [e[: self.matryoshka_dim] for e in embed] if self.matryoshka_dim else embed
         )
@@ -119,9 +97,9 @@ class CustomEmbedding(BaseModel, Embeddings):
             List[float]: The embedded query vector.
         """
         try:
-            embed = self.hosted_embedding.embed_query(text)
-        except Exception as e:
-            logging.warning(f"Issue with hosted embedding, moving to CPU: {e}")
             embed = self.cpu_embedding.embed_query(text)
+        except Exception as e:
+            logging.warning(f"Issue with CPU embedding: {e}")
+            embed = []
         logging.warning(text)
         return embed[: self.matryoshka_dim] if self.matryoshka_dim else embed