add target_label_id

app.py

@@ -25,6 +25,10 @@ def main():
         'Выберите языковую модель',
         ('sberbank-ai/rugpt3small_based_on_gpt2',)
     )
+    cls_model_config = transformers.AutoConfig.from_pretrained(cls_model_name)
+    label2id = cls_model_config.label2id
+    label_key = st.selectbox("Веберите нужный атрибут текста", label2id.keys())
+    target_label_id = label2id[label_key]
     prompt = st.text_input("Начало текста:", "Привет")
     alpha = st.slider("Alpha:", min_value=-10, max_value=10, step=1)
     auth_token = os.environ.get('TOKEN') or True
@@ -50,7 +54,9 @@ def load_sampler(cls_model_name, lm_tokenizer):
 
 
 @st.cache
-def inference(lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool = True, alpha: float = 5) -> str:
+def inference(
+        lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool = True, alpha: float = 5, target_label_id: int = 0
+) -> str:
     generator = load_generator(lm_model_name=lm_model_name)
     lm_tokenizer = transformers.AutoTokenizer.from_pretrained(lm_model_name)
     caif_sampler = load_sampler(cls_model_name=cls_model_name, lm_tokenizer=lm_tokenizer)
@@ -61,6 +67,7 @@ def inference(lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool =
         "temperature": 1.0,
         "top_k_classifier": 100,
         "classifier_weight": alpha,
+        "target_cls_id": target_label_id
     }
     generator.set_ordinary_sampler(ordinary_sampler)
     if device == "cpu":
@@ -74,7 +81,6 @@ def inference(lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool =
             input_prompt=prompt,
             max_length=20,
             caif_period=1,
-            caif_tokens_num=100,
             entropy=None,
             **kwargs
         )

sampling.py

@@ -49,10 +49,11 @@ class CAIFSampler:
         top_k_classifier,
         classifier_weight,
         caif_tokens_num=None,
+        act_type: str = "softmax",
         **kwargs
     ):
+        target_cls_id = kwargs["target_cls_id"]
         next_token_logits = output_logis[:, -1]
-
         next_token_log_probs = F.log_softmax(
             next_token_logits, dim=-1
         )
@@ -63,7 +64,8 @@ class CAIFSampler:
             temperature,
             top_k_classifier,
             classifier_weight,
-            caif_tokens_num=caif_tokens_num
+            caif_tokens_num=caif_tokens_num,
+            target_cls_id=target_cls_id
         )
         topk_probs = next_token_unnormalized_probs.topk(top_k, -1)
         next_tokens = sample_from_values(
@@ -80,6 +82,7 @@ class CAIFSampler:
         temperature,
         top_k_classifier,
         classifier_weight,
+        target_cls_id: int = 0,
         caif_tokens_num=None
     ):
 
@@ -109,7 +112,7 @@ class CAIFSampler:
             )
         else:
             classifier_log_probs = self.get_classifier_log_probs(
-                classifier_input, caif_tokens_num=caif_tokens_num
+                classifier_input, caif_tokens_num=caif_tokens_num, target_cls_id=target_cls_id,
             ).view(-1, top_k_classifier)
 
         next_token_probs = torch.exp(
@@ -118,7 +121,7 @@ class CAIFSampler:
         )
         return next_token_probs, top_next_token_log_probs[1]
 
-    def get_classifier_log_probs(self, input, caif_tokens_num=None):
+    def get_classifier_log_probs(self, input, caif_tokens_num=None, target_cls_id: int = 0):
         input_ids = self.classifier_tokenizer(
             input, padding=True, return_tensors="pt"
         ).to(self.device)
@@ -128,7 +131,7 @@ class CAIFSampler:
                 input_ids["attention_mask"] = input_ids["attention_mask"][:, -caif_tokens_num:]
             if "token_type_ids" in input_ids.keys():
                 input_ids["token_type_ids"] = input_ids["token_type_ids"][:, -caif_tokens_num:]
-        logits = self.classifier_model(**input_ids).logits[:, 0].squeeze(-1)
+        logits = self.classifier_model(**input_ids).logits[:, target_cls_id].squeeze(-1)
         return torch.log(torch.sigmoid(logits))
 
     def get_classifier_probs(self, input, caif_tokens_num=None):