File size: 1,076 Bytes
14a87b0 3133b5e 14a87b0 3133b5e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
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]
) -> PreTrainedModel:
from adapters import AutoAdapterModel
model = AutoAdapterModel.from_pretrained(**model_kwargs)
model.load_adapter(set_active=True, **adapter_kwargs)
return model
|