| 3 | |
| 4 | |
| 5 | class PromptSpell(torch.nn.Module): |
| 6 | def __init__(self, spell_length, hidden_size, spell_func): |
| 7 | super(PromptSpell, self).__init__() |
| 8 | self.spell_length = spell_length |
| 9 | self.hidden_size = hidden_size |
| 10 | self.spell_embeddings = torch.nn.Embedding(self.spell_length, self.hidden_size) |
| 11 | self.spell_func = spell_func |
| 12 | if self.spell_func == "lstm": |
| 13 | self.lstm_head = torch.nn.LSTM(input_size=self.hidden_size, |
| 14 | hidden_size=self.hidden_size, |
| 15 | num_layers=2, |
| 16 | # dropout=self.lstm_dropout, |
| 17 | bidirectional=True, |
| 18 | batch_first=True) # .to(torch.device("cuda")) |
| 19 | self.mlp_head = torch.nn.Sequential(torch.nn.Linear(2 * self.hidden_size, self.hidden_size), |
| 20 | torch.nn.ReLU(), |
| 21 | torch.nn.Linear(self.hidden_size, self.hidden_size)) |
| 22 | elif self.spell_func == "mlp": |
| 23 | self.mlp_head = torch.nn.Sequential(torch.nn.Linear(self.hidden_size, self.hidden_size), |
| 24 | torch.nn.ReLU(), |
| 25 | torch.nn.Linear(self.hidden_size, self.hidden_size)) |
| 26 | elif self.spell_func != "none": |
| 27 | raise NotImplementedError("Prompt function " + self.spell_func) |
| 28 | |
| 29 | def init_embedding(self, word_embeddings=None, task_tokens=None): |
| 30 | num_words = 5000 |
| 31 | with torch.no_grad(): |
| 32 | for i in range(self.spell_length): |
| 33 | rand_token = random.randrange(num_words) |
| 34 | if task_tokens is None: |
| 35 | target_embedding = word_embeddings[rand_token] |
| 36 | else: |
| 37 | word_embedding = word_embeddings[rand_token] |
| 38 | task_token = random.choice(task_tokens) |
| 39 | task_embedding = word_embeddings[task_token] |
| 40 | ratio = random.random() |
| 41 | target_embedding = word_embedding * ratio + task_embedding * (1 - ratio) |
| 42 | self.spell_embeddings.weight.data[i] = target_embedding |
| 43 | |
| 44 | def forward(self): |
| 45 | prompt_embeds = self.spell_embeddings.weight.unsqueeze(0) |
| 46 | if self.spell_func == "lstm": |
| 47 | prompt_embeds = self.lstm_head(prompt_embeds)[0] |
| 48 | if self.spell_func == "lstm" or self.spell_func == "mlp": |
| 49 | prompt_embeds = self.mlp_head(prompt_embeds) |
| 50 | return prompt_embeds |