fix load_model_with_adapter location

configs/retriever/related_span_retriever_with_relations_from_other_docs.yaml

@@ -18,7 +18,7 @@ vectorstore:
   embedding:
     _target_: src.langchain_modules.HuggingFaceSpanEmbeddings
     model:
-      _target_: src.models.utils.adapters.load_model_with_adapter
+      _target_: src.models.utils.load_model_with_adapter
       model_kwargs:
         pretrained_model_name_or_path: allenai/specter2_base
       adapter_kwargs:

src/models/utils/__init__.py

@@ -0,0 +1 @@
+from .loading import load_model_from_pie_model, load_model_with_adapter, load_tokenizer_from_pie_taskmodule

src/models/utils/loading.py

@@ -0,0 +1,32 @@
+from typing import Any, Dict
+
+from pie_modules.models import *  # noqa: F403
+from pie_modules.taskmodules import *  # noqa: F403
+from pytorch_ie import AutoModel, AutoTaskModule, PyTorchIEModel, TaskModule
+from pytorch_ie.models import *  # noqa: F403
+from pytorch_ie.taskmodules import *  # noqa: F403
+from transformers import PreTrainedModel, PreTrainedTokenizer
+
+
+def load_model_from_pie_model(model_kwargs: Dict[str, Any]) -> PreTrainedModel:
+
+    pie_model: PyTorchIEModel = AutoModel.from_pretrained(**model_kwargs)
+
+    return pie_model.model.model
+
+
+def load_tokenizer_from_pie_taskmodule(taskmodule_kwargs: Dict[str, Any]) -> PreTrainedTokenizer:
+
+    pie_taskmodule: TaskModule = AutoTaskModule.from_pretrained(**taskmodule_kwargs)
+
+    return pie_taskmodule.tokenizer
+
+
+def load_model_with_adapter(
+        model_kwargs: Dict[str, Any], adapter_kwargs: Dict[str, Any]
+) -> "ModelAdaptersMixin":
+    from adapters import AutoAdapterModel, ModelAdaptersMixin
+
+    model = AutoAdapterModel.from_pretrained(**model_kwargs)
+    model.load_adapter(set_active=True, **adapter_kwargs)
+    return model