| |
| import os |
| from pathlib import Path |
| import torch |
| import torch.nn as nn |
| from transformers import BertConfig, BertModel, BertForSequenceClassification |
| from safetensors.torch import load_file |
| from huggingface_hub import hf_hub_download |
|
|
| |
| _orig_named_params = nn.Module.named_parameters |
| def _compat_named_parameters(self, prefix: str = '', recurse: bool = True, **kwargs): |
| return _orig_named_params(self, prefix=prefix, recurse=recurse) |
| nn.Module.named_parameters = _compat_named_parameters |
|
|
| class NeuralOpsBertConfig(BertConfig): |
| model_type = "bert" |
|
|
| def _load_attention_weights_into(model, repo_or_dir): |
| p = Path(repo_or_dir) |
| if p.is_dir() and (p / "attention_weights.safetensors").exists(): |
| w_path = str(p / "attention_weights.safetensors") |
| else: |
| w_path = hf_hub_download(str(repo_or_dir), "attention_weights.safetensors") |
| weights = load_file(w_path) |
| model.load_state_dict(weights, strict=False) |
|
|
| class NeuralOpsBertModel(BertModel): |
| config_class = NeuralOpsBertConfig |
| def __init__(self, config, add_pooling_layer=True): |
| super().__init__(config, add_pooling_layer=add_pooling_layer) |
| from neuralops.api import replace_attention |
| replace_attention(self, attention="phantom", copy_weights=False) |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): |
| config = kwargs.pop("config", None) or cls.config_class.from_pretrained(pretrained_model_name_or_path) |
| base = BertModel.from_pretrained("bert-base-uncased", config=config, **kwargs) |
| from neuralops.api import replace_attention |
| replace_attention(base, attention="phantom", copy_weights=False) |
| _load_attention_weights_into(base, pretrained_model_name_or_path) |
| return base |
|
|
| class NeuralOpsBertForSequenceClassification(BertForSequenceClassification): |
| config_class = NeuralOpsBertConfig |
| def __init__(self, config): |
| super().__init__(config) |
| from neuralops.api import replace_attention |
| replace_attention(self, attention="phantom", copy_weights=False) |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): |
| config = kwargs.pop("config", None) or cls.config_class.from_pretrained(pretrained_model_name_or_path) |
| base = BertForSequenceClassification.from_pretrained("bert-base-uncased", config=config, **kwargs) |
| from neuralops.api import replace_attention |
| replace_attention(base, attention="phantom", copy_weights=False) |
| _load_attention_weights_into(base, pretrained_model_name_or_path) |
| return base |
|
|
