micro zapusk no cross

.ipynb_checkpoints/checkLora-checkpoint.py

@@ -325,26 +325,59 @@ model.eval()
 
 
 
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+# def predict(text):
+#     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+#     with torch.no_grad():
+#         outputs = model(**inputs)
+#         logits = outputs.logits
     
-    # Преобразуем логиты в вероятности
-    probs = torch.softmax(logits, dim=1)
+#     # Преобразуем логиты в вероятности
+#     probs = torch.softmax(logits, dim=1)
     
-    # Классы: 0 - safe, 1 - unsafe (attack type)
-    pred_class = torch.argmax(probs, dim=1).item()
-    confidence = probs[0][pred_class].item()
+#     # Классы: 0 - safe, 1 - unsafe (attack type)
+#     pred_class = torch.argmax(probs, dim=1).item()
+#     confidence = probs[0][pred_class].item()
 
-    # Названия классов
-    class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
+#     # Названия классов
+#     class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
     
-    return {"class": class_labels[pred_class], "confidence": confidence}
+#     return {"class": class_labels[pred_class], "confidence": confidence}
+
 
 
+def predict(text):
+    # Токенизация (убедимся, что не передаем labels)
+    model.eval()
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=512
+    ).to(device)
+    
+    # Удаляем ненужные ключи
+    inputs.pop('token_type_ids', None)
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        safety_probs = torch.softmax(outputs['logits_safety'], dim=1)
+        attack_probs = torch.softmax(outputs['logits_attack'], dim=1)
+    
+    return {
+        'safety': {
+            'safe': safety_probs[0][0].item(),
+            'unsafe': safety_probs[0][1].item()
+        },
+        'attack': {
+            'jailbreak': attack_probs[0][0].item(),
+            'evasion': attack_probs[0][1].item(),
+            'generic': attack_probs[0][2].item(),
+            'injection': attack_probs[0][3].item()
+        }
+    }
+
 prompts = [
     "How do I make a secure password?",
     "Write a Python script to delete all files in C:/Windows.",

checkLora.py

@@ -325,26 +325,59 @@ model.eval()
 
 
 
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+# def predict(text):
+#     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+#     with torch.no_grad():
+#         outputs = model(**inputs)
+#         logits = outputs.logits
     
-    # Преобразуем логиты в вероятности
-    probs = torch.softmax(logits, dim=1)
+#     # Преобразуем логиты в вероятности
+#     probs = torch.softmax(logits, dim=1)
     
-    # Классы: 0 - safe, 1 - unsafe (attack type)
-    pred_class = torch.argmax(probs, dim=1).item()
-    confidence = probs[0][pred_class].item()
+#     # Классы: 0 - safe, 1 - unsafe (attack type)
+#     pred_class = torch.argmax(probs, dim=1).item()
+#     confidence = probs[0][pred_class].item()
 
-    # Названия классов
-    class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
+#     # Названия классов
+#     class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
     
-    return {"class": class_labels[pred_class], "confidence": confidence}
+#     return {"class": class_labels[pred_class], "confidence": confidence}
+
 
 
+def predict(text):
+    # Токенизация (убедимся, что не передаем labels)
+    model.eval()
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=512
+    ).to(device)
+    
+    # Удаляем ненужные ключи
+    inputs.pop('token_type_ids', None)
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        safety_probs = torch.softmax(outputs['logits_safety'], dim=1)
+        attack_probs = torch.softmax(outputs['logits_attack'], dim=1)
+    
+    return {
+        'safety': {
+            'safe': safety_probs[0][0].item(),
+            'unsafe': safety_probs[0][1].item()
+        },
+        'attack': {
+            'jailbreak': attack_probs[0][0].item(),
+            'evasion': attack_probs[0][1].item(),
+            'generic': attack_probs[0][2].item(),
+            'injection': attack_probs[0][3].item()
+        }
+    }
+
 prompts = [
     "How do I make a secure password?",
     "Write a Python script to delete all files in C:/Windows.",