Upload bge_custom_impl.py with huggingface_hub

bge_custom_impl.py

@@ -1,3 +1,4 @@
+from calendar import c
 import json
 import os
 from typing import Any
@@ -7,7 +8,7 @@ from transformers import AutoModel, AutoTokenizer
 
 class ConcatCustomPooling(nn.Module):
     def __init__(self, model_name_or_path="BAAI/bge-large-en-v1.5",layers=None,max_seq_len=512, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__()
         self.layers = layers
         self.base_name = model_name_or_path
         self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
@@ -47,7 +48,8 @@ class ConcatCustomPooling(nn.Module):
         with open(os.path.join(save_dir, "config.json"), "w") as fOut:
             json.dump(self.get_config_dict(), fOut, indent=4)
 
-    def load(self,load_dir: str, **kwargs) -> "ConcatCustomPooling":
+    @classmethod
+    def load(cls,load_dir: str, **kwargs) -> "ConcatCustomPooling":
         with open(os.path.join(load_dir, "config.json")) as fIn:
             config = json.load(fIn)