micro zapusk no cross

.ipynb_checkpoints/checkLora-checkpoint.py

@@ -143,7 +143,7 @@ import torch
 
 import torch
 from peft import PeftModel, PeftConfig
-from transformers import BertTokenizer, AutoModel
+from transformers import BertTokenizer, BertForSequenceClassification
 
 # Определяем устройство
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -151,8 +151,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # Загружаем конфигурацию LoRA
 config = PeftConfig.from_pretrained("./micro_no_cross_fine_tuned")
 
-# Загружаем базовую модель
-model = AutoModel.from_pretrained(config.base_model_name_or_path)
+# Загружаем базовую модель для классификации с двумя выходами (безопасность и атака)
+model = BertForSequenceClassification.from_pretrained(config.base_model_name_or_path, num_labels=5)  # 1 для безопасности и 4 для атак
 
 # Загружаем LoRA-адаптеры
 model = PeftModel.from_pretrained(model, "./micro_no_cross_fine_tuned", is_trainable=False)
@@ -171,12 +171,12 @@ def predict(text):
 
     # Включаем режим без градиентов для предсказания
     with torch.no_grad():
-        # Используем ваш кастомный forward метод (MultiTaskBert)
+        # Получаем выходы модели
         outputs = model(**inputs)
-        
-        # Достаем logits из outputs
-        safety_probs = torch.softmax(outputs['logits_safety'], dim=1)
-        attack_probs = torch.softmax(outputs['logits_attack'], dim=1)
+
+        # Получаем вероятности для безопасности
+        safety_probs = torch.softmax(outputs.logits[:, 0:2], dim=1)  # Пример для двух классов: безопасно/небезопасно
+        attack_probs = torch.softmax(outputs.logits[:, 2:], dim=1)  # Пример для атак: jailbreak, evasion, generic, injection
 
     return {
         'safety': {
@@ -202,6 +202,7 @@ for prompt in test_prompts:
     result = predict(prompt)
     print(f"\nТекст: {prompt}")
     print(f"Безопасность: Safe {result['safety']['safe']:.2%} | Unsafe {result['safety']['unsafe']:.2%}")
+    
     if result['safety']['unsafe'] > 0.5:
         print("Тип атаки:")
         for attack_type, prob in result['attack'].items():

checkLora.py

@@ -143,7 +143,7 @@ import torch
 
 import torch
 from peft import PeftModel, PeftConfig
-from transformers import BertTokenizer, AutoModel
+from transformers import BertTokenizer, BertForSequenceClassification
 
 # Определяем устройство
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -151,8 +151,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # Загружаем конфигурацию LoRA
 config = PeftConfig.from_pretrained("./micro_no_cross_fine_tuned")
 
-# Загружаем базовую модель
-model = AutoModel.from_pretrained(config.base_model_name_or_path)
+# Загружаем базовую модель для классификации с двумя выходами (безопасность и атака)
+model = BertForSequenceClassification.from_pretrained(config.base_model_name_or_path, num_labels=5)  # 1 для безопасности и 4 для атак
 
 # Загружаем LoRA-адаптеры
 model = PeftModel.from_pretrained(model, "./micro_no_cross_fine_tuned", is_trainable=False)
@@ -171,12 +171,12 @@ def predict(text):
 
     # Включаем режим без градиентов для предсказания
     with torch.no_grad():
-        # Используем ваш кастомный forward метод (MultiTaskBert)
+        # Получаем выходы модели
         outputs = model(**inputs)
-        
-        # Достаем logits из outputs
-        safety_probs = torch.softmax(outputs['logits_safety'], dim=1)
-        attack_probs = torch.softmax(outputs['logits_attack'], dim=1)
+
+        # Получаем вероятности для безопасности
+        safety_probs = torch.softmax(outputs.logits[:, 0:2], dim=1)  # Пример для двух классов: безопасно/небезопасно
+        attack_probs = torch.softmax(outputs.logits[:, 2:], dim=1)  # Пример для атак: jailbreak, evasion, generic, injection
 
     return {
         'safety': {
@@ -202,6 +202,7 @@ for prompt in test_prompts:
     result = predict(prompt)
     print(f"\nТекст: {prompt}")
     print(f"Безопасность: Safe {result['safety']['safe']:.2%} | Unsafe {result['safety']['unsafe']:.2%}")
+    
     if result['safety']['unsafe'] > 0.5:
         print("Тип атаки:")
         for attack_type, prob in result['attack'].items():