wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-10 12:20:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

add util for getting the base model

Browse Source
This commit is contained in:
Sourab Mangrulkar
2023-02-22 00:14:24 +05:30
parent e6ef85a711
commit 1ef0f89a0c
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/peft/peft_model.py
+7 -4
View File
@@ -284,10 +284,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
"""
Forward pass of the model.
"""
if isinstance(self.peft_config, PromptLearningConfig):
return self.base_model(*args, **kwargs)
else:
return self.base_model.model(*args, **kwargs)
return self.get_base_model()(*args, **kwargs)
@contextmanager
def disable_adapter(self):
@@ -305,6 +302,12 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
else:
self.base_model.enable_adapter_layers()
def get_base_model(self):
"""
Returns the base model.
"""
return self.base_model if isinstance(self.peft_config, PromptLearningConfig) else self.base_model.model
class PeftModelForSequenceClassification(PeftModel):
"""