chore: use binary classifier

app.py

@@ -19,10 +19,10 @@ chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, dtype=torch.b
 chat_model.to("cuda")
 chat_tokenizer = AutoTokenizer.from_pretrained(chat_model_name)
 
-moderator_model_name = "saiteki-kai/QA-DeBERTa-v3-large"
+moderator_model_name = "saiteki-kai/QA-DeBERTa-v3-large-binary-3"
 moderator_model = AutoModelForSequenceClassification.from_pretrained(moderator_model_name, device_map="cpu")
 moderator_model.to("cuda")
-moderator_tokenizer = AutoTokenizer.from_pretrained(moderator_model_name)
+moderator_tokenizer = AutoTokenizer.from_pretrained(moderator_model_name, padding_side="right")
 
 def generate_responses(model, tokenizer, prompts):
     messages = [[{"role": "user", "content": message}] for message in prompts]
@@ -51,14 +51,13 @@ def classify_pairs(model, tokenizer, prompts, responses):
         prompt + "[SEP]" + response for prompt, response in zip(prompts, responses)
     ]
 
-    input_ids = tokenizer(texts, padding=True, truncation=True, max_length=1024, return_tensors="pt").to(model.device)
+    input_ids = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(model.device)
+    print(tokenizer.batch_decode(input_ids["input_ids"]))
 
     with torch.inference_mode():
         outputs = model(**input_ids)
-        scores = torch.sigmoid(outputs.logits).detach().cpu()
-        scores = [[float(score) for score in s if float(score) > 0.5] for s in scores]
-        
-        unsafety_scores = [sum(s) / len(s) if len(s) > 0 else 0.0 for s in scores]
+        scores = torch.softmax(outputs.logits, dim=-1).detach().cpu()
+        unsafety_scores = [float(s[1]) for s in scores] # get unsafe axis
 
     return unsafety_scores