Update embedding_model.py

embedding_model.py

@@ -6,11 +6,11 @@ import torch
 from sentence_transformers import models
 from transformers import AutoModel
 
-MODEL_PATH_IN_MODULE = "0_Model"
-
 
 class EmbeddingModel(models.Transformer):
     def __init__(self, *args, **kwargs):
+        self.model_name_or_path = "lamarr-llm-development/elbedding"
+        kwargs.pop("model_name_or_path", None)
         super().__init__(*args, **kwargs)
 
     def tokenize(
@@ -63,15 +63,6 @@ class EmbeddingModel(models.Transformer):
         with open(os.path.join(save_dir, "sentence_bert_config.json"), "w+") as f:
             json.dump(self.get_config_dict(), f, indent=4)
 
-        model_path = os.path.join(save_dir, MODEL_PATH_IN_MODULE)
-        if not os.path.exists(model_path):
-            os.makedirs(model_path)
-
-        with open(f"{model_path}/config.json", "w+") as f:
-            json.dump(self.get_config_dict(), f, indent=4)
-
     @staticmethod
-    def load(load_dir: str, **kwargs) -> "EmbeddingModel":
-        with open(os.path.join(load_dir, "config.json")) as fIn:
-            config = json.load(fIn)
-        return EmbeddingModel(**config, **kwargs)
+    def load(**kwargs) -> "EmbeddingModel":
+        return EmbeddingModel(**kwargs)