feat: add sigmoid

app.py

@@ -19,12 +19,10 @@ moderator_tokenizer = AutoTokenizer.from_pretrained(moderator_model_name)
 def generate_responses(model, tokenizer, prompts):
     messages = [[{"role": "user", "content": message}] for message in prompts]
 
-    texts = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
+    texts = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(model.device)
 
     with torch.inference_mode():
-        model_inputs = tokenizer(texts, padding=True, return_tensors="pt").to(model.device)
         generated_ids = model.generate(
             **model_inputs,
             do_sample=False,
@@ -47,11 +45,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)
+
     with torch.inference_mode():
-        input_ids = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(model.device)
-        outputs = model(**input_ids).logits
+        outputs = model(**input_ids)
+        scores = torch.sigmoid(outputs.logits)
 
-    return outputs
+    return scores
 
 
 @spaces.GPU()