Spaces:
Running
Running
File size: 1,764 Bytes
e0b11c9 d75dc6d e0b11c9 d75dc6d e0b11c9 |
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 34 35 36 37 38 39 40 41 42 43 |
import torch
from transformers import pipeline
from .model_utils import Hack_no_grad
from .steers import Projected_Adaptor
from .model_base import LMSteerBase
class Switching_GPTNeoModel(LMSteerBase):
def __init__(self, model_name, adapted_component, adaptor_class,
num_steers, rank, epsilon, init_var,
low_resource_mode):
super().__init__()
self.adapted_component = adapted_component
self.pipeline = pipeline('text-generation', model=model_name)
self.model = self.pipeline.model
self.tokenizer = self.pipeline.tokenizer
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
self.init_var = init_var
self.num_steers = num_steers
self.device = torch.device("cpu")
embed_dim = self.model.lm_head.weight.shape[1]
vocab_size = self.model.lm_head.weight.shape[0]
for _param in self.model.parameters():
_param.requires_grad_(False)
if adapted_component == "final_layer":
self.model.transformer = Hack_no_grad(self.model.transformer)
self.steer = Projected_Adaptor(
self.model.lm_head, adaptor_class, num_steers, embed_dim,
vocab_size, rank, epsilon, init_var, "output")
self.model.set_output_embeddings(self.steer)
elif adapted_component == "input_embedding":
self.steer = Projected_Adaptor(
self.model.transformer.wte, adaptor_class, num_steers,
embed_dim, vocab_size, rank, epsilon, init_var, "input")
self.model.transformer.set_input_embeddings(self.steer)
else:
raise NotImplementedError()
|