fix

app.py

@@ -24,14 +24,11 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 ATTRIBUTE_MODELS = {
     "Russian": (
         "cointegrated/rubert-tiny-toxicity",
-        'tinkoff-ai/response-quality-classifier-tiny',
-        'tinkoff-ai/response-quality-classifier-base',
-        'tinkoff-ai/response-quality-classifier-large',
-        "SkolkovoInstitute/roberta_toxicity_classifier",
         "SkolkovoInstitute/russian_toxicity_classifier"
     ),
     "English": (
         "unitary/toxic-bert",
+        "distilbert-base-uncased-finetuned-sst-2-english"
     )
 }
 
@@ -72,7 +69,7 @@ WARNING_TEXT = {
     "English": """
     **Warning!**
     
-    If you are clicking checkbox bellow positive""" + r"$\alpha$" + """ values for CAIF sampling become available.
+    If you are clicking checkbox bellow positive """ + r"$\alpha$" + """ values for CAIF sampling become available.
     It means that language model will be forced to produce toxic or/and abusive text.
     This space is only a demonstration of our method for controllable text generation 
     and we are not responsible for the content produced by this method.
@@ -128,11 +125,17 @@ def main():
         label2id = cls_model_config.label2id
         label_key = st.selectbox(ATTRIBUTE_LABEL[language], label2id.keys())
         target_label_id = label2id[label_key]
-    else:
+        act_type = "sigmoid"
+    elif cls_model_config.problem_type == "single_label_classification":
         label2id = cls_model_config.label2id
-        print(list(label2id.keys()))
         label_key = st.selectbox(ATTRIBUTE_LABEL[language], [list(label2id.keys())[-1]])
         target_label_id = 1
+        act_type = "sigmoid"
+    else:
+        label2id = cls_model_config.label2id
+        label_key = st.selectbox(ATTRIBUTE_LABEL[language], label2id.keys())
+        target_label_id = label2id[label_key]
+        act_type = "softmax"
     st.write(WARNING_TEXT[language])
     show_pos_alpha = st.checkbox("Show positive alphas", value=False)
     prompt = st.text_input(TEXT_PROMPT_LABEL[language], PROMPT_EXAMPLE[language])
@@ -168,6 +171,7 @@ def main():
             target_label_id=target_label_id,
             entropy_threshold=entropy_threshold,
             fp16=fp16,
+            act_type=act_type
         )
     st.subheader("Generated text:")
     st.write(text)

sampling.py

@@ -53,6 +53,7 @@ class CAIFSampler:
         **kwargs
     ):
         target_cls_id = kwargs["target_cls_id"]
+        act_type = kwargs["act_type"]
         next_token_logits = output_logis[:, -1]
         next_token_log_probs = F.log_softmax(
             next_token_logits, dim=-1
@@ -83,6 +84,7 @@ class CAIFSampler:
         top_k_classifier,
         classifier_weight,
         target_cls_id: int = 0,
+        act_type: str = "sigmoid",
         caif_tokens_num=None
     ):
 
@@ -107,12 +109,15 @@ class CAIFSampler:
         if self.invert_cls_probs:
             classifier_log_probs = torch.log(
                 1 - self.get_classifier_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)
             )
         else:
             classifier_log_probs = self.get_classifier_log_probs(
-                classifier_input, caif_tokens_num=caif_tokens_num, target_cls_id=target_cls_id,
+                classifier_input,
+                caif_tokens_num=caif_tokens_num,
+                target_cls_id=target_cls_id,
+                act_type=act_type,
             ).view(-1, top_k_classifier)
 
         next_token_probs = torch.exp(
@@ -121,7 +126,7 @@ class CAIFSampler:
         )
         return next_token_probs, top_next_token_log_probs[1]
 
-    def get_classifier_log_probs(self, input, caif_tokens_num=None, target_cls_id: int = 0):
+    def get_classifier_log_probs(self, input, caif_tokens_num=None, target_cls_id: int = 0, act_type: str = "sigmoid"):
         input_ids = self.classifier_tokenizer(
             input, padding=True, return_tensors="pt"
         ).to(self.device)
@@ -131,10 +136,15 @@ 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[:, target_cls_id].squeeze(-1)
-        return torch.log(torch.sigmoid(logits))
 
-    def get_classifier_probs(self, input, caif_tokens_num=None):
+        if act_type == "sigmoid":
+            logits = self.classifier_model(**input_ids).logits[:, target_cls_id].squeeze(-1)
+            return F.logsigmoid(logits)
+        if act_type == "softmax":
+            logits = F.log_softmax(self.classifier_model(**input_ids).logits)[:, target_cls_id].squeeze(-1)
+            return logits
+
+    def get_classifier_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)
@@ -142,5 +152,5 @@ class CAIFSampler:
             input_ids["input_ids"] = input_ids["input_ids"][-caif_tokens_num:]
             if "attention_mask" in input_ids.keys():
                 input_ids["attention_mask"] = input_ids["attention_mask"][-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.sigmoid(logits)