Update pseudo_moe_st_module.py

pseudo_moe_st_module.py

@@ -5,6 +5,10 @@ from transformers import AutoModel, AutoTokenizer, AutoConfig
 
 
 class PseudoMoETransformer(nn.Module):
+    # Save this module in the model's root folder (as in this repo) so that the copied pseudo_moe_st_module.py is
+    # found again on reload; otherwise Sentence Transformers saves it in a subfolder and reloading fails.
+    save_in_root = True
+
     def __init__(
         self,
         model_name_or_path: str = None,
@@ -51,7 +55,8 @@ class PseudoMoETransformer(nn.Module):
 
         model_kwargs = kwargs.get("model_kwargs", {})
         config_kwargs = kwargs.get("config_kwargs", {})
-        tokenizer_kwargs = kwargs.get("tokenizer_kwargs", {})
+        # Sentence Transformers 6 passes tokenizer options as processor_kwargs (formerly tokenizer_kwargs); accept both.
+        tokenizer_kwargs = kwargs.get("processor_kwargs") or kwargs.get("tokenizer_kwargs") or {}
 
         return PseudoMoETransformer(
             model_name_or_path=load_path,
@@ -84,9 +89,10 @@ class PseudoMoETransformer(nn.Module):
                 )
             domain = self.default_domain
         else:
-            if domain not in self.config.domain_names:
+            # "general" (the base model without LoRA) is always valid but is not listed in config.domain_names.
+            if domain != "general" and domain not in self.config.domain_names:
                 raise ValueError(
-                    f"Invalid domain: {domain}. Must be one of {self.config.domain_names}."
+                    f"Invalid domain: {domain}. Must be one of {['general'] + list(self.config.domain_names)}."
                 )
         self.set_domain(domain)
         
@@ -130,4 +136,4 @@ class PseudoMoETransformer(nn.Module):
 
     @property
     def domain(self) -> str:
-        return self.auto_model.domain
+        return self.auto_model.domain