mirror of
https://github.com/wassname/peft.git
synced 2026-10-30 13:21:29 +08:00
add util for getting the base model
This commit is contained in:
1 parent
e6ef85a711
commit
1ef0f89a0c
1 file changed
+7
-4
@@ -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):
|
||||
"""
|
||||
|
||||
Reference in new issue
Block a user