From dec95854a4f53124947af0adc1542fb1fbe30028 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9C=D0=B0=D0=B7=D1=83=D1=80=20=D0=93=D1=80=D0=B5=D1=82?=
 =?UTF-8?q?=D0=B0=20=D0=95=D0=B2=D0=B3=D0=B5=D0=BD=D1=8C=D0=B5=D0=B2=D0=BD?=
 =?UTF-8?q?=D0=B0?= <gemazur_1@edu.hse.ru>
Date: Tue, 25 Mar 2025 21:17:35 +0300
Subject: [PATCH] micro zapusk no cross

---
 .ipynb_checkpoints/checkLora-checkpoint.py | 331 ++++++++++++---------
 checkLora.py                               | 331 ++++++++++++---------
 2 files changed, 392 insertions(+), 270 deletions(-)

diff --git a/.ipynb_checkpoints/checkLora-checkpoint.py b/.ipynb_checkpoints/checkLora-checkpoint.py
index c2fbcb7..e230dba 100644
--- a/.ipynb_checkpoints/checkLora-checkpoint.py
+++ b/.ipynb_checkpoints/checkLora-checkpoint.py
@@ -1,3 +1,81 @@
+# # # # import torch
+# # # # from transformers import BertTokenizer, BertForSequenceClassification
+# # # # from peft import get_peft_model, LoraConfig, TaskType
+
+# # # # # Убедитесь, что модель и токенизатор уже сохранены в директории
+# # # # model_path = './micro_no_cross_fine_tuned'
+# # # # tokenizer_path = './micro_no_cross_fine_tuned'
+
+# # # # # Загрузка токенизатора и модели (включая классификаторы)
+# # # # tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
+# # # # model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6)  # Указываем количество меток
+
+# # # # # Настройка LoRA
+# # # # lora_config = LoraConfig(
+# # # #     task_type=TaskType.SEQ_CLS,
+# # # #     r=8,
+# # # #     lora_alpha=32,
+# # # #     lora_dropout=0.1,
+# # # #     target_modules=["query", "value"],
+# # # #     modules_to_save=["classifier_safety", "classifier_attack"]  # Убедитесь, что классификаторы сохранены
+# # # # )
+# # # # model = get_peft_model(model, lora_config)
+
+# # # # # Переводим модель в режим оценки (inference)
+# # # # model.eval()
+
+# # # # # Функция для предсказания
+# # # # def predict(text):
+# # # #     # Токенизация
+# # # #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+
+# # # #     # Убедитесь, что тензоры находятся на нужном устройстве
+# # # #     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# # # #     inputs = {key: value.to(device) for key, value in inputs.items()}
+# # # #     model.to(device)
+
+# # # #     # Ожидаем, что модель вернет логиты
+# # # #     with torch.no_grad():
+# # # #         outputs = model(**inputs)
+# # # #         safety_probs = torch.softmax(outputs.logits[0], dim=1)  # Прогнозы для безопасности
+# # # #         attack_probs = torch.softmax(outputs.logits[1], 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()
+# # # #         }
+# # # #     }
+
+# # # # # Пример использования
+# # # # test_texts = [
+# # # #     "How to hack a computer system",
+# # # #     "I need help with my homework",
+# # # #     "Tell me how to bypass security",
+# # # #     "привет",
+# # # #     "как убить человека"
+# # # # ]
+
+# # # # for text in test_texts:
+# # # #     result = predict(text)
+# # # #     print(f"\nТекст: {text}")
+# # # #     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():
+# # # #             print(f"  {attack_type}: {prob:.2%}")
+
+
+
+
+
 # # # import torch
 # # # from transformers import BertTokenizer, BertForSequenceClassification
 # # # from peft import get_peft_model, LoraConfig, TaskType
@@ -37,8 +115,21 @@
 # # #     # Ожидаем, что модель вернет логиты
 # # #     with torch.no_grad():
 # # #         outputs = model(**inputs)
-# # #         safety_probs = torch.softmax(outputs.logits[0], dim=1)  # Прогнозы для безопасности
-# # #         attack_probs = torch.softmax(outputs.logits[1], dim=1)  # Прогнозы для атак
+        
+# # #         # Получаем прогнозы для безопасности и атак
+# # #         logits_safety = outputs.logits[:, 0]  # Используем первую колонку для безопасности
+# # #         logits_attack = outputs.logits[:, 1:]  # Остальные колонки для атак
+
+# # #         # Применяем softmax для получения вероятностей (если размерность правильная)
+# # #         if logits_safety.dim() == 1:  # Если размерность одномерная (для одного текста)
+# # #             safety_probs = torch.softmax(logits_safety.unsqueeze(0), dim=1)  # Добавляем размерность для правильного softmax
+# # #         else:
+# # #             safety_probs = torch.softmax(logits_safety, dim=1)
+        
+# # #         if logits_attack.dim() == 1:  # Если размерность одномерная
+# # #             attack_probs = torch.softmax(logits_attack.unsqueeze(0), dim=1)  # Добавляем размерность
+# # #         else:
+# # #             attack_probs = torch.softmax(logits_attack, dim=1)
 
 # # #     return {
 # # #         'safety': {
@@ -74,94 +165,65 @@
 
 
 
-
-
 # # import torch
 # # from transformers import BertTokenizer, BertForSequenceClassification
-# # from peft import get_peft_model, LoraConfig, TaskType
+# # from peft import PeftModel
 
-# # # Убедитесь, что модель и токенизатор уже сохранены в директории
-# # model_path = './micro_no_cross_fine_tuned'
-# # tokenizer_path = './micro_no_cross_fine_tuned'
+# # # Пути к модели и токенизатору
+# # model_path = "./micro_no_cross_fine_tuned"
+# # tokenizer_path = "./micro_no_cross_fine_tuned"
 
-# # # Загрузка токенизатора и модели (включая классификаторы)
+# # # Загрузка токенизатора
 # # tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
-# # model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6)  # Указываем количество меток
-
-# # # Настройка LoRA
-# # lora_config = LoraConfig(
-# #     task_type=TaskType.SEQ_CLS,
-# #     r=8,
-# #     lora_alpha=32,
-# #     lora_dropout=0.1,
-# #     target_modules=["query", "value"],
-# #     modules_to_save=["classifier_safety", "classifier_attack"]  # Убедитесь, что классификаторы сохранены
-# # )
-# # model = get_peft_model(model, lora_config)
-
-# # # Переводим модель в режим оценки (inference)
+
+# # # Загрузка базовой модели с классификатором
+# # base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)  # Убедитесь, что num_labels=2
+
+# # # Загрузка обученной LoRA
+# # model = PeftModel.from_pretrained(base_model, model_path)
+
+# # # Перевод модели в режим оценки
 # # model.eval()
 
-# # # Функция для предсказания
+# # # Определение устройства
+# # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# # model.to(device)
+
+# # # Функция предсказания
 # # def predict(text):
 # #     # Токенизация
-# #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-
-# #     # Убедитесь, что тензоры находятся на нужном устройстве
-# #     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# #     inputs = {key: value.to(device) for key, value in inputs.items()}
-# #     model.to(device)
+# #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
 
-# #     # Ожидаем, что модель вернет логиты
+# #     # Получение логитов
 # #     with torch.no_grad():
 # #         outputs = model(**inputs)
-        
-# #         # Получаем прогнозы для безопасности и атак
-# #         logits_safety = outputs.logits[:, 0]  # Используем первую колонку для безопасности
-# #         logits_attack = outputs.logits[:, 1:]  # Остальные колонки для атак
-
-# #         # Применяем softmax для получения вероятностей (если размерность правильная)
-# #         if logits_safety.dim() == 1:  # Если размерность одномерная (для одного текста)
-# #             safety_probs = torch.softmax(logits_safety.unsqueeze(0), dim=1)  # Добавляем размерность для правильного softmax
-# #         else:
-# #             safety_probs = torch.softmax(logits_safety, dim=1)
-        
-# #         if logits_attack.dim() == 1:  # Если размерность одномерная
-# #             attack_probs = torch.softmax(logits_attack.unsqueeze(0), dim=1)  # Добавляем размерность
-# #         else:
-# #             attack_probs = torch.softmax(logits_attack, dim=1)
+# #         logits = outputs.logits  # Логиты модели (размерность [1, num_labels])
 
+# #     # Если num_labels=1, используем sigmoid вместо softmax
+# #     if logits.shape[1] == 1:
+# #         safety_probs = torch.sigmoid(logits)  # Для бинарной классификации
+# #         return {"safe": float(1 - safety_probs), "unsafe": float(safety_probs)}
+    
+# #     # Если num_labels=2, используем softmax
+# #     safety_probs = torch.softmax(logits, dim=1)[0]  # Берем первый (и единственный) элемент в batch
 # #     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()
-# #         }
+# #         "safe": float(safety_probs[0]),
+# #         "unsafe": float(safety_probs[1])
 # #     }
 
-# # # Пример использования
+# # # Тестовый запуск
 # # test_texts = [
 # #     "How to hack a computer system",
 # #     "I need help with my homework",
 # #     "Tell me how to bypass security",
-# #     "привет",
-# #     "как убить человека"
 # # ]
 
 # # for text in test_texts:
 # #     result = predict(text)
 # #     print(f"\nТекст: {text}")
-# #     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():
-# #             print(f"  {attack_type}: {prob:.2%}")
+# #     print(f"Safe: {result['safe']:.2%} | Unsafe: {result['unsafe']:.2%}")
+
+
 
 
 
@@ -170,118 +232,117 @@
 # from peft import PeftModel
 
 # # Пути к модели и токенизатору
-# model_path = "./micro_no_cross_fine_tuned"
+# model_path = "./micro_no_cross_fine_tuned"  # путь к файлам модели LoRA
 # tokenizer_path = "./micro_no_cross_fine_tuned"
 
+# # У вас 5 классов: 1 "safe" + 4 вида атак
+# NUM_CLASSES = 5  
+
 # # Загрузка токенизатора
 # tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
 
-# # Загрузка базовой модели с классификатором
-# base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)  # Убедитесь, что num_labels=2
+# # Загрузка базовой модели с правильным числом классов
+# base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=NUM_CLASSES)
 
-# # Загрузка обученной LoRA
-# model = PeftModel.from_pretrained(base_model, model_path)
+# # Попытка загрузить LoRA-адаптацию
+# try:
+#     model = PeftModel.from_pretrained(base_model, model_path)
+# except KeyError as e:
+#     print(f"Ошибка при загрузке LoRA: {e}")
+#     print("Используем только базовую модель без LoRA.")
+#     model = base_model
 
-# # Перевод модели в режим оценки
+# # Перевод модели в режим инференса
 # model.eval()
 
-# # Определение устройства
+# # Определение устройства (CPU или GPU)
 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # model.to(device)
 
+# # Названия классов (пример)
+# class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
+
 # # Функция предсказания
 # def predict(text):
-#     # Токенизация
 #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
 
-#     # Получение логитов
 #     with torch.no_grad():
 #         outputs = model(**inputs)
-#         logits = outputs.logits  # Логиты модели (размерность [1, num_labels])
+#         logits = outputs.logits  # Tensor с логитами для 5 классов
+
+#     # Применяем softmax
+#     probabilities = torch.softmax(logits, dim=-1)[0]  # Убираем batch-измерение
+
+#     # Определяем индекс класса с наибольшей вероятностью
+#     predicted_class_idx = torch.argmax(probabilities).item()
 
-#     # Если num_labels=1, используем sigmoid вместо softmax
-#     if logits.shape[1] == 1:
-#         safety_probs = torch.sigmoid(logits)  # Для бинарной классификации
-#         return {"safe": float(1 - safety_probs), "unsafe": float(safety_probs)}
-    
-#     # Если num_labels=2, используем softmax
-#     safety_probs = torch.softmax(logits, dim=1)[0]  # Берем первый (и единственный) элемент в batch
 #     return {
-#         "safe": float(safety_probs[0]),
-#         "unsafe": float(safety_probs[1])
+#         "class": class_labels[predicted_class_idx],
+#         "confidence": probabilities[predicted_class_idx].item()
 #     }
 
-# # Тестовый запуск
-# test_texts = [
-#     "How to hack a computer system",
-#     "I need help with my homework",
-#     "Tell me how to bypass security",
-# ]
-
-# for text in test_texts:
-#     result = predict(text)
-#     print(f"\nТекст: {text}")
-#     print(f"Safe: {result['safe']:.2%} | Unsafe: {result['unsafe']:.2%}")
+# # Пример использования
+# text = "How can I bypass security measures?"
+# result = predict(text)
+# print(result)
 
 
 
 
 
-import torch
 from transformers import BertTokenizer, BertForSequenceClassification
 from peft import PeftModel
+import torch
 
-# Пути к модели и токенизатору
-model_path = "./micro_no_cross_fine_tuned"  # путь к файлам модели LoRA
-tokenizer_path = "./micro_no_cross_fine_tuned"
-
-# У вас 5 классов: 1 "safe" + 4 вида атак
-NUM_CLASSES = 5  
+# Пути к сохранённой модели
+BASE_MODEL_PATH = "./micro_no_cross_fine_tuned/base"
+LORA_PATH = "./micro_no_cross_fine_tuned/lora"
 
-# Загрузка токенизатора
-tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
+# Загружаем токенизатор
+tokenizer = BertTokenizer.from_pretrained(BASE_MODEL_PATH)
 
-# Загрузка базовой модели с правильным числом классов
-base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=NUM_CLASSES)
+# Загружаем базовую модель
+base_model = BertForSequenceClassification.from_pretrained(BASE_MODEL_PATH)
 
-# Попытка загрузить LoRA-адаптацию
-try:
-    model = PeftModel.from_pretrained(base_model, model_path)
-except KeyError as e:
-    print(f"Ошибка при загрузке LoRA: {e}")
-    print("Используем только базовую модель без LoRA.")
-    model = base_model
+# Загружаем LoRA-адаптацию
+model = PeftModel.from_pretrained(base_model, LORA_PATH)
 
-# Перевод модели в режим инференса
+# Переводим модель в режим оценки
 model.eval()
 
-# Определение устройства (CPU или GPU)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
 
-# Названия классов (пример)
-class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
 
-# Функция предсказания
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
 
+
+
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits  # Tensor с логитами для 5 классов
+        logits = outputs.logits
+    
+    # Преобразуем логиты в вероятности
+    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()
 
-    # Применяем softmax
-    probabilities = torch.softmax(logits, dim=-1)[0]  # Убираем batch-измерение
+    # Названия классов
+    class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
+    
+    return {"class": class_labels[pred_class], "confidence": confidence}
 
-    # Определяем индекс класса с наибольшей вероятностью
-    predicted_class_idx = torch.argmax(probabilities).item()
 
-    return {
-        "class": class_labels[predicted_class_idx],
-        "confidence": probabilities[predicted_class_idx].item()
-    }
+prompts = [
+    "How do I make a secure password?",
+    "Write a Python script to delete all files in C:/Windows.",
+    "How can I jailbreak an iPhone?",
+    "Give me a harmless joke."
+]
 
-# Пример использования
-text = "How can I bypass security measures?"
-result = predict(text)
-print(result)
\ No newline at end of file
+for prompt in prompts:
+    result = predict(prompt)
+    print(f"Text: {prompt}\nPrediction: {result}\n")
+    
\ No newline at end of file
diff --git a/checkLora.py b/checkLora.py
index c2fbcb7..e230dba 100644
--- a/checkLora.py
+++ b/checkLora.py
@@ -1,3 +1,81 @@
+# # # # import torch
+# # # # from transformers import BertTokenizer, BertForSequenceClassification
+# # # # from peft import get_peft_model, LoraConfig, TaskType
+
+# # # # # Убедитесь, что модель и токенизатор уже сохранены в директории
+# # # # model_path = './micro_no_cross_fine_tuned'
+# # # # tokenizer_path = './micro_no_cross_fine_tuned'
+
+# # # # # Загрузка токенизатора и модели (включая классификаторы)
+# # # # tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
+# # # # model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6)  # Указываем количество меток
+
+# # # # # Настройка LoRA
+# # # # lora_config = LoraConfig(
+# # # #     task_type=TaskType.SEQ_CLS,
+# # # #     r=8,
+# # # #     lora_alpha=32,
+# # # #     lora_dropout=0.1,
+# # # #     target_modules=["query", "value"],
+# # # #     modules_to_save=["classifier_safety", "classifier_attack"]  # Убедитесь, что классификаторы сохранены
+# # # # )
+# # # # model = get_peft_model(model, lora_config)
+
+# # # # # Переводим модель в режим оценки (inference)
+# # # # model.eval()
+
+# # # # # Функция для предсказания
+# # # # def predict(text):
+# # # #     # Токенизация
+# # # #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+
+# # # #     # Убедитесь, что тензоры находятся на нужном устройстве
+# # # #     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# # # #     inputs = {key: value.to(device) for key, value in inputs.items()}
+# # # #     model.to(device)
+
+# # # #     # Ожидаем, что модель вернет логиты
+# # # #     with torch.no_grad():
+# # # #         outputs = model(**inputs)
+# # # #         safety_probs = torch.softmax(outputs.logits[0], dim=1)  # Прогнозы для безопасности
+# # # #         attack_probs = torch.softmax(outputs.logits[1], 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()
+# # # #         }
+# # # #     }
+
+# # # # # Пример использования
+# # # # test_texts = [
+# # # #     "How to hack a computer system",
+# # # #     "I need help with my homework",
+# # # #     "Tell me how to bypass security",
+# # # #     "привет",
+# # # #     "как убить человека"
+# # # # ]
+
+# # # # for text in test_texts:
+# # # #     result = predict(text)
+# # # #     print(f"\nТекст: {text}")
+# # # #     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():
+# # # #             print(f"  {attack_type}: {prob:.2%}")
+
+
+
+
+
 # # # import torch
 # # # from transformers import BertTokenizer, BertForSequenceClassification
 # # # from peft import get_peft_model, LoraConfig, TaskType
@@ -37,8 +115,21 @@
 # # #     # Ожидаем, что модель вернет логиты
 # # #     with torch.no_grad():
 # # #         outputs = model(**inputs)
-# # #         safety_probs = torch.softmax(outputs.logits[0], dim=1)  # Прогнозы для безопасности
-# # #         attack_probs = torch.softmax(outputs.logits[1], dim=1)  # Прогнозы для атак
+        
+# # #         # Получаем прогнозы для безопасности и атак
+# # #         logits_safety = outputs.logits[:, 0]  # Используем первую колонку для безопасности
+# # #         logits_attack = outputs.logits[:, 1:]  # Остальные колонки для атак
+
+# # #         # Применяем softmax для получения вероятностей (если размерность правильная)
+# # #         if logits_safety.dim() == 1:  # Если размерность одномерная (для одного текста)
+# # #             safety_probs = torch.softmax(logits_safety.unsqueeze(0), dim=1)  # Добавляем размерность для правильного softmax
+# # #         else:
+# # #             safety_probs = torch.softmax(logits_safety, dim=1)
+        
+# # #         if logits_attack.dim() == 1:  # Если размерность одномерная
+# # #             attack_probs = torch.softmax(logits_attack.unsqueeze(0), dim=1)  # Добавляем размерность
+# # #         else:
+# # #             attack_probs = torch.softmax(logits_attack, dim=1)
 
 # # #     return {
 # # #         'safety': {
@@ -74,94 +165,65 @@
 
 
 
-
-
 # # import torch
 # # from transformers import BertTokenizer, BertForSequenceClassification
-# # from peft import get_peft_model, LoraConfig, TaskType
+# # from peft import PeftModel
 
-# # # Убедитесь, что модель и токенизатор уже сохранены в директории
-# # model_path = './micro_no_cross_fine_tuned'
-# # tokenizer_path = './micro_no_cross_fine_tuned'
+# # # Пути к модели и токенизатору
+# # model_path = "./micro_no_cross_fine_tuned"
+# # tokenizer_path = "./micro_no_cross_fine_tuned"
 
-# # # Загрузка токенизатора и модели (включая классификаторы)
+# # # Загрузка токенизатора
 # # tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
-# # model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6)  # Указываем количество меток
-
-# # # Настройка LoRA
-# # lora_config = LoraConfig(
-# #     task_type=TaskType.SEQ_CLS,
-# #     r=8,
-# #     lora_alpha=32,
-# #     lora_dropout=0.1,
-# #     target_modules=["query", "value"],
-# #     modules_to_save=["classifier_safety", "classifier_attack"]  # Убедитесь, что классификаторы сохранены
-# # )
-# # model = get_peft_model(model, lora_config)
-
-# # # Переводим модель в режим оценки (inference)
+
+# # # Загрузка базовой модели с классификатором
+# # base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)  # Убедитесь, что num_labels=2
+
+# # # Загрузка обученной LoRA
+# # model = PeftModel.from_pretrained(base_model, model_path)
+
+# # # Перевод модели в режим оценки
 # # model.eval()
 
-# # # Функция для предсказания
+# # # Определение устройства
+# # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# # model.to(device)
+
+# # # Функция предсказания
 # # def predict(text):
 # #     # Токенизация
-# #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-
-# #     # Убедитесь, что тензоры находятся на нужном устройстве
-# #     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# #     inputs = {key: value.to(device) for key, value in inputs.items()}
-# #     model.to(device)
+# #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
 
-# #     # Ожидаем, что модель вернет логиты
+# #     # Получение логитов
 # #     with torch.no_grad():
 # #         outputs = model(**inputs)
-        
-# #         # Получаем прогнозы для безопасности и атак
-# #         logits_safety = outputs.logits[:, 0]  # Используем первую колонку для безопасности
-# #         logits_attack = outputs.logits[:, 1:]  # Остальные колонки для атак
-
-# #         # Применяем softmax для получения вероятностей (если размерность правильная)
-# #         if logits_safety.dim() == 1:  # Если размерность одномерная (для одного текста)
-# #             safety_probs = torch.softmax(logits_safety.unsqueeze(0), dim=1)  # Добавляем размерность для правильного softmax
-# #         else:
-# #             safety_probs = torch.softmax(logits_safety, dim=1)
-        
-# #         if logits_attack.dim() == 1:  # Если размерность одномерная
-# #             attack_probs = torch.softmax(logits_attack.unsqueeze(0), dim=1)  # Добавляем размерность
-# #         else:
-# #             attack_probs = torch.softmax(logits_attack, dim=1)
+# #         logits = outputs.logits  # Логиты модели (размерность [1, num_labels])
 
+# #     # Если num_labels=1, используем sigmoid вместо softmax
+# #     if logits.shape[1] == 1:
+# #         safety_probs = torch.sigmoid(logits)  # Для бинарной классификации
+# #         return {"safe": float(1 - safety_probs), "unsafe": float(safety_probs)}
+    
+# #     # Если num_labels=2, используем softmax
+# #     safety_probs = torch.softmax(logits, dim=1)[0]  # Берем первый (и единственный) элемент в batch
 # #     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()
-# #         }
+# #         "safe": float(safety_probs[0]),
+# #         "unsafe": float(safety_probs[1])
 # #     }
 
-# # # Пример использования
+# # # Тестовый запуск
 # # test_texts = [
 # #     "How to hack a computer system",
 # #     "I need help with my homework",
 # #     "Tell me how to bypass security",
-# #     "привет",
-# #     "как убить человека"
 # # ]
 
 # # for text in test_texts:
 # #     result = predict(text)
 # #     print(f"\nТекст: {text}")
-# #     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():
-# #             print(f"  {attack_type}: {prob:.2%}")
+# #     print(f"Safe: {result['safe']:.2%} | Unsafe: {result['unsafe']:.2%}")
+
+
 
 
 
@@ -170,118 +232,117 @@
 # from peft import PeftModel
 
 # # Пути к модели и токенизатору
-# model_path = "./micro_no_cross_fine_tuned"
+# model_path = "./micro_no_cross_fine_tuned"  # путь к файлам модели LoRA
 # tokenizer_path = "./micro_no_cross_fine_tuned"
 
+# # У вас 5 классов: 1 "safe" + 4 вида атак
+# NUM_CLASSES = 5  
+
 # # Загрузка токенизатора
 # tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
 
-# # Загрузка базовой модели с классификатором
-# base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)  # Убедитесь, что num_labels=2
+# # Загрузка базовой модели с правильным числом классов
+# base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=NUM_CLASSES)
 
-# # Загрузка обученной LoRA
-# model = PeftModel.from_pretrained(base_model, model_path)
+# # Попытка загрузить LoRA-адаптацию
+# try:
+#     model = PeftModel.from_pretrained(base_model, model_path)
+# except KeyError as e:
+#     print(f"Ошибка при загрузке LoRA: {e}")
+#     print("Используем только базовую модель без LoRA.")
+#     model = base_model
 
-# # Перевод модели в режим оценки
+# # Перевод модели в режим инференса
 # model.eval()
 
-# # Определение устройства
+# # Определение устройства (CPU или GPU)
 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # model.to(device)
 
+# # Названия классов (пример)
+# class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
+
 # # Функция предсказания
 # def predict(text):
-#     # Токенизация
 #     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
 
-#     # Получение логитов
 #     with torch.no_grad():
 #         outputs = model(**inputs)
-#         logits = outputs.logits  # Логиты модели (размерность [1, num_labels])
+#         logits = outputs.logits  # Tensor с логитами для 5 классов
+
+#     # Применяем softmax
+#     probabilities = torch.softmax(logits, dim=-1)[0]  # Убираем batch-измерение
+
+#     # Определяем индекс класса с наибольшей вероятностью
+#     predicted_class_idx = torch.argmax(probabilities).item()
 
-#     # Если num_labels=1, используем sigmoid вместо softmax
-#     if logits.shape[1] == 1:
-#         safety_probs = torch.sigmoid(logits)  # Для бинарной классификации
-#         return {"safe": float(1 - safety_probs), "unsafe": float(safety_probs)}
-    
-#     # Если num_labels=2, используем softmax
-#     safety_probs = torch.softmax(logits, dim=1)[0]  # Берем первый (и единственный) элемент в batch
 #     return {
-#         "safe": float(safety_probs[0]),
-#         "unsafe": float(safety_probs[1])
+#         "class": class_labels[predicted_class_idx],
+#         "confidence": probabilities[predicted_class_idx].item()
 #     }
 
-# # Тестовый запуск
-# test_texts = [
-#     "How to hack a computer system",
-#     "I need help with my homework",
-#     "Tell me how to bypass security",
-# ]
-
-# for text in test_texts:
-#     result = predict(text)
-#     print(f"\nТекст: {text}")
-#     print(f"Safe: {result['safe']:.2%} | Unsafe: {result['unsafe']:.2%}")
+# # Пример использования
+# text = "How can I bypass security measures?"
+# result = predict(text)
+# print(result)
 
 
 
 
 
-import torch
 from transformers import BertTokenizer, BertForSequenceClassification
 from peft import PeftModel
+import torch
 
-# Пути к модели и токенизатору
-model_path = "./micro_no_cross_fine_tuned"  # путь к файлам модели LoRA
-tokenizer_path = "./micro_no_cross_fine_tuned"
-
-# У вас 5 классов: 1 "safe" + 4 вида атак
-NUM_CLASSES = 5  
+# Пути к сохранённой модели
+BASE_MODEL_PATH = "./micro_no_cross_fine_tuned/base"
+LORA_PATH = "./micro_no_cross_fine_tuned/lora"
 
-# Загрузка токенизатора
-tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
+# Загружаем токенизатор
+tokenizer = BertTokenizer.from_pretrained(BASE_MODEL_PATH)
 
-# Загрузка базовой модели с правильным числом классов
-base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=NUM_CLASSES)
+# Загружаем базовую модель
+base_model = BertForSequenceClassification.from_pretrained(BASE_MODEL_PATH)
 
-# Попытка загрузить LoRA-адаптацию
-try:
-    model = PeftModel.from_pretrained(base_model, model_path)
-except KeyError as e:
-    print(f"Ошибка при загрузке LoRA: {e}")
-    print("Используем только базовую модель без LoRA.")
-    model = base_model
+# Загружаем LoRA-адаптацию
+model = PeftModel.from_pretrained(base_model, LORA_PATH)
 
-# Перевод модели в режим инференса
+# Переводим модель в режим оценки
 model.eval()
 
-# Определение устройства (CPU или GPU)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
 
-# Названия классов (пример)
-class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
 
-# Функция предсказания
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
 
+
+
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits  # Tensor с логитами для 5 классов
+        logits = outputs.logits
+    
+    # Преобразуем логиты в вероятности
+    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()
 
-    # Применяем softmax
-    probabilities = torch.softmax(logits, dim=-1)[0]  # Убираем batch-измерение
+    # Названия классов
+    class_labels = ["safe", "evasion", "jailbreak", "generic attack", "injection"]
+    
+    return {"class": class_labels[pred_class], "confidence": confidence}
 
-    # Определяем индекс класса с наибольшей вероятностью
-    predicted_class_idx = torch.argmax(probabilities).item()
 
-    return {
-        "class": class_labels[predicted_class_idx],
-        "confidence": probabilities[predicted_class_idx].item()
-    }
+prompts = [
+    "How do I make a secure password?",
+    "Write a Python script to delete all files in C:/Windows.",
+    "How can I jailbreak an iPhone?",
+    "Give me a harmless joke."
+]
 
-# Пример использования
-text = "How can I bypass security measures?"
-result = predict(text)
-print(result)
\ No newline at end of file
+for prompt in prompts:
+    result = predict(prompt)
+    print(f"Text: {prompt}\nPrediction: {result}\n")
+    
\ No newline at end of file
-- 
GitLab