Upload prompt.py

prompt.py

@@ -0,0 +1,112 @@
+from transformers import AutoModelForMaskedLM , AutoTokenizer
+import torch
+
+class Prompting(object):
+  """ doc string 
+   This class helps us to implement
+   Prompt-based Learning Model
+  """
+  def __init__(self, **kwargs):
+    """ constructor 
+
+    parameter:
+    ----------
+       model: AutoModelForMaskedLM
+            path to a Pre-trained language model form HuggingFace Hub
+       tokenizer: AutoTokenizer
+            path to tokenizer if different tokenizer is used, 
+            otherwise leave it empty
+    """
+    model_path=kwargs['model']
+    tokenizer_path= kwargs['model']
+    if "tokenizer" in kwargs.keys():
+      tokenizer_path= kwargs['tokenizer']
+    self.model = AutoModelForMaskedLM.from_pretrained(model_path)
+    self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+  def prompt_pred(self,text):
+    """
+      Predict MASK token by listing the probability of candidate tokens 
+      where the first token is the most likely
+
+      Parameters:
+      ----------
+      text: str 
+          The text including [MASK] token.
+          It supports single MASK token. If more [MASK]ed tokens 
+          are given, it takes the first one.
+
+      Returns:
+      --------
+      list of (token, prob)
+         The return is a list of all token in LM Vocab along with 
+         their prob score, sort by score in descending order 
+    """
+    indexed_tokens=self.tokenizer(text, return_tensors="pt").input_ids
+    tokenized_text= self.tokenizer.convert_ids_to_tokens (indexed_tokens[0])
+    # take the first masked token
+    mask_pos=tokenized_text.index(self.tokenizer.mask_token)
+    self.model.eval()
+    with torch.no_grad():
+      outputs = self.model(indexed_tokens)
+      predictions = outputs[0]
+    values, indices=torch.sort(predictions[0, mask_pos],  descending=True)
+    #values=torch.nn.functional.softmax(values, dim=0)
+    result=list(zip(self.tokenizer.convert_ids_to_tokens(indices), values))
+    self.scores_dict={a:b for a,b in result}
+    return result
+
+  def compute_tokens_prob(self, text, token_list1, token_list2):
+    """
+    Compute the activations for given two token list, 
+
+    Parameters:
+    ---------
+    token_list1: List(str)
+     it is a list for positive polarity tokens such as good, great. 
+    token_list2: List(str)
+     it is a list for negative polarity tokens such as bad, terrible.      
+
+    Returns:
+    --------
+    Tuple (
+       the probability for first token list,
+       the probability of the second token list,
+       the ratio score1/ (score1+score2)
+       The softmax returns
+    """
+    _=self.prompt_pred(text)
+    score1=[self.scores_dict[token1] if token1 in self.scores_dict.keys() else 0\
+            for token1 in token_list1]
+    score1= sum(score1)
+    score2=[self.scores_dict[token2] if token2 in self.scores_dict.keys() else 0\
+            for token2 in token_list2]
+    score2= sum(score2)
+    softmax_rt=torch.nn.functional.softmax(torch.Tensor([score1,score2]), dim=0)
+    return softmax_rt
+
+  def fine_tune(self, sentences, labels, prompt=" Since it was [MASK].",goodToken="good",badToken="bad"):
+    """  
+      Fine tune the model
+    """
+    good=tokenizer.convert_tokens_to_ids(goodToken)
+    bad=tokenizer.convert_tokens_to_ids(badToken)
+
+    from transformers import AdamW
+    optimizer = AdamW(self.model.parameters(),lr=1e-3)
+
+    for sen, label in zip(sentences, labels):
+      tokenized_text = self.tokenizer.tokenize(sen+prompt)
+      indexed_tokens = self.tokenizer.convert_tokens_to_ids(tokenized_text)
+      tokens_tensor = torch.tensor([indexed_tokens])
+      # take the first masked token
+      mask_pos=tokenized_text.index(self.tokenizer.mask_token)
+      outputs = self.model(tokens_tensor)
+      predictions = outputs[0]
+      pred=predictions[0, mask_pos][[good,bad]]
+      prob=torch.nn.functional.softmax(pred, dim=0)
+      lossFunc = torch.nn.CrossEntropyLoss()
+      loss=lossFunc(prob.unsqueeze(0), torch.tensor([label]))
+      loss.backward()
+      optimizer.step()
+    print("done!")