Update splade.py

splade.py

@@ -45,20 +45,26 @@ class Splade(PreTrainedModel):
     def __init__(self, config, weights_path=None):
         super().__init__(config)
         self.name = "splade"
+    
         base_cfg = AutoConfig.from_pretrained(
             config.model_name_or_path,
             attn_implementation=config.attn_implementation,
             torch_dtype="auto",
         )
+    
         self.tokenizer = prepare_tokenizer(
             config.model_name_or_path, padding_side=config.padding_side
         )
+    
         if is_flash_attn_2_available():
             config.attn_implementation = "flash_attention_2"
         else:
             config.attn_implementation = "sdpa"
+
+        source = weights_path or config.model_name_or_path
+    
         self.model = get_decoder_model(
-            model_name_or_path=config.model_name_or_path,
+            model_name_or_path=source,
             attn_implementation=config.attn_implementation,
             bidirectional=getattr(config, "bidirectional", False),
             base_cfg=base_cfg,