Update README.md

README.md

@@ -273,7 +273,44 @@ class LevelToScorePipeline(TextGenerationPipeline):
         return records
 
 
-from src.rank_dicts import SingleLabelRankDict
+class SingleLabelRankDict:
+    def __init__(
+        self,
+        rank_dict: Dict[Text, Any]
+    ):
+        self._rank_dict = rank_dict
+        
+    def __len__(self) -> int:
+        return len(self._rank_dict)
+        
+    def get_rank_dict(self, tokenizer: PreTrainedTokenizer) -> Dict[int, Any]:
+        return {tokenizer.convert_tokens_to_ids([token])[0]: value for token, value in self._rank_dict.items()}
+    
+    def to_tokenizer(self, tokenizer: PreTrainedTokenizer) -> PreTrainedTokenizer:
+        """Augment tokenizer vocab with `rank_dict` IN-PLACE.
+        """
+        vocabs: List[Text] = self._rank_dict.keys()
+        new_vocab = [vocab for vocab in vocabs if vocab not in tokenizer.get_vocab()]
+        tokenizer.add_tokens(new_vocab)
+        return tokenizer
+        def from_tokenizer(cls, tokenizer: PreTrainedTokenizer) -> "SingleLabelRankDict":
+        vocab = tokenizer.get_vocab()
+        rank_dict = {}
+        pattern = re.compile(r" <\|label_level_(\d+)\|>")
+        
+        for token in vocab.keys():
+            match = pattern.match(token)
+            if match:
+                value = int(match.group(1))
+                # normalized_value = value / (len(vocab) - 1)
+                rank_dict[token] = value
+                
+        # normalize rank_values
+        num_levels = max(rank_dict.values()) + 1
+        for token in rank_dict.keys():
+            rank_dict[token] = 1. / num_levels * (rank_dict[token] + 0.5)
+        
+        return cls(rank_dict=rank_dict)
 
 
 model = transformers.AutoModelForCausalLM.from_pretrained(
@@ -327,17 +364,17 @@ premise = "Sam is sleeping."
 hypothesis = "Sam is awake."
 
 inputs = [
-  {
-    "role": "user",
-    "content": "### Question: Given the premise \"{premise}\", how likely is it that the hypothesis \"{hypothesis}\" is true?\n\n".format(
-      premise=premise,
-      hypothesis=hypothesis
-    )
-  },
-  {
-    "role": "assitant",
-    "content": "### Answer:"
-  }
+    {
+        "role": "user",
+        "content": "### Question: Given the premise \"{premise}\", how likely is it that the hypothesis \"{hypothesis}\" is true?\n\n".format(
+            premise=premise,
+            hypothesis=hypothesis
+        )
+    },
+    {
+        "role": "assitant",
+        "content": "### Answer:"
+    }
 ]
     
 result = pipe(inputs)