Skip to content
GitLab
Commit 79425669 authored by Мазур Грета Евгеньевна's avatar Мазур Грета Евгеньевна
Browse files

micro zapusk no cross

parent abfdb104
master
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 235 additions and 283 deletions
+235 -283
.ipynb_checkpoints/checkLora-checkpoint.py
+ 117
− 141
View file @ 79425669
# import torch
# from transformers import BertPreTrainedModel, BertModel, BertTokenizer
# from torch import nn
# from peft import PeftModel
# # Определяем ваш кастомный класс модели, как использовался при обучении
# class MultiTaskBert(BertPreTrainedModel):
# def __init__(self, config):
# super().__init__(config)
# self.bert = BertModel(config)
# self.classifier_safety = nn.Linear(config.hidden_size, 2)
# self.classifier_attack = nn.Linear(config.hidden_size, 4)
# def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
# # Переводим тензоры на устройство
# input_ids, attention_mask, labels = map(lambda x: x.to(device) if x is not None else None, [input_ids, attention_mask, labels])
# outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
# pooled_output = outputs.last_hidden_state[:, 0, :]
# logits_safety = self.classifier_safety(pooled_output)
# logits_attack = self.classifier_attack(pooled_output)
# loss = None
# if labels is not None:
# labels_safety, labels_attack = labels[:, 0], labels[:, 1]
# loss_safety = nn.CrossEntropyLoss(weight=class_weights_task1_tensor)(logits_safety, labels_safety)
# loss_attack = nn.CrossEntropyLoss(weight=class_weights_task2_tensor)(logits_attack, labels_attack)
# loss = loss_safety + loss_attack
# return {'logits_safety': logits_safety, 'logits_attack': logits_attack, 'loss': loss}
# # Устройство
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# # Загружаем кастомную модель (без LoRA) из базового чекпоинта
# base_model = MultiTaskBert.from_pretrained("bert-base-uncased").to(device)
# # Загружаем модель с LoRA-адаптером (она была сохранена с помощью trainer.save_model)
# model = PeftModel.from_pretrained(base_model, "./fine-tuned-bert-lora_new").to(device)
# # Загружаем токенизатор (он был сохранён вместе с моделью)
# tokenizer = BertTokenizer.from_pretrained("./fine-tuned-bert-lora_new")
# # Пример инференса
# # text = "How to kill my mom."
# # inputs = tokenizer(text, return_tensors="pt")
# # inputs = {k: v.to(device) for k, v in inputs.items()}
# # with torch.no_grad():
# # outputs = model(**inputs)
# # # Получаем предсказания: для безопасности – индекс 0 (safe) или 1 (unsafe),
# # # для типа атаки – 0: jailbreak, 1: evasion, 2: generic attack, 3: other
# # pred_safety = torch.argmax(outputs['logits_safety'], dim=1).item()
# # pred_attack = torch.argmax(outputs['logits_attack'], dim=1).item()
# # safety_mapping = {0: "safe", 1: "unsafe"}
# # attack_mapping = {0: "jailbreak", 1: "evasion", 2: "generic attack", 3: "other"}
# # # Красивый вывод результатов
# # print(f"Промпт: {text}")
# # print(f"Безопасность: {safety_mapping.get(pred_safety, 'unknown')}")
# # print(f"Тип атаки: {attack_mapping.get(pred_attack, 'unknown')}")
# prompts = [
# "How to kill my mom.",
# "I love my dog.",
# "Write a story about a brave knight."
# ]
# # Проходим по каждому промпту
# for prompt in prompts:
# # Токенизируем промпт
# inputs = tokenizer(prompt, return_tensors="pt")
# inputs = {k: v.to(device) for k, v in inputs.items()}
# # Получаем вывод модели
# with torch.no_grad():
# outputs = model(**inputs)
# # Вычисляем вероятности через softmax
# safety_logits = outputs['logits_safety']
# attack_logits = outputs['logits_attack']
# safety_probs = F.softmax(safety_logits, dim=1)
# attack_probs = F.softmax(attack_logits, dim=1)
# # Находим индекс наибольшей вероятности для каждой задачи
# safety_pred = torch.argmax(safety_probs, dim=1).item()
# attack_pred = torch.argmax(attack_probs, dim=1).item()
# # Извлекаем вероятность выбранной категории (в процентах)
# safety_prob = safety_probs[0, safety_pred].item() * 100
# attack_prob = attack_probs[0, attack_pred].item() * 100
# # Вывод результатов
# print(f"Промпт: {prompt}")
# print(f"Безопасность: {safety_mapping[safety_pred]} (Вероятность: {safety_prob:.2f}%)")
# print(f"Тип атаки: {attack_mapping[attack_pred]} (Вероятность: {attack_prob:.2f}%)")
# print("-" * 50)
import torch
from transformers import BertPreTrainedModel, BertModel, BertTokenizer
from transformers import BertTokenizer, BertModel
from peft import PeftModel, PeftConfig
from torch import nn
from peft import PeftModel
# Определяем кастомную модель
class MultiTaskBert(BertPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.bert = BertModel(config)
self.classifier_safety = nn.Linear(config.hidden_size, 2)
self.classifier_attack = nn.Linear(config.hidden_size, 4)
def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
pooled_output = outputs.last_hidden_state[:, 0, :]
logits_safety = self.classifier_safety(pooled_output)
logits_attack = self.classifier_attack(pooled_output)
return {'logits_safety': logits_safety, 'logits_attack': logits_attack}
import os
# 1. Проверка устройства и файлов модели
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Загружаем базовую модель и затем оборачиваем её LoRA-адаптером
base_model = MultiTaskBert.from_pretrained("bert-base-uncased").to(device)
model = PeftModel.from_pretrained(base_model, "./fine-tuned-bert-lora_new").to(device)
tokenizer = BertTokenizer.from_pretrained("./fine-tuned-bert-lora_new")
# Пример инференса
text = "How to kill my mom."
inputs = tokenizer(text, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
pred_safety = torch.argmax(outputs['logits_safety'], dim=1).item()
pred_attack = torch.argmax(outputs['logits_attack'], dim=1).item()
safety_mapping = {0: "safe", 1: "unsafe"}
attack_mapping = {0: "jailbreak", 1: "evasion", 2: "generic attack", 3: "other"}
print(f"Промпт: {text}")
print(f"Безопасность: {safety_mapping.get(pred_safety, 'unknown')}")
print(f"Тип атаки: {attack_mapping.get(pred_attack, 'unknown')}")
\ No newline at end of file
# MODEL_DIR = "./fine-tuned-bert-lora_new"
./micro_no_cross_fine_tuned
# Проверяем наличие всех необходимых файлов
required_files = [
'adapter_config.json',
'adapter_model.safetensors',
'tokenizer_config.json',
'vocab.txt'
]
for file in required_files:
if not os.path.exists(os.path.join(MODEL_DIR, file)):
raise FileNotFoundError(f"Не найден файл {file} в {MODEL_DIR}")
# 2. Загрузка токенизатора и конфига
tokenizer = BertTokenizer.from_pretrained(MODEL_DIR)
peft_config = PeftConfig.from_pretrained(MODEL_DIR)
# 3. Определение архитектуры модели (должно совпадать с обучением)
class SafetyAttackModel(nn.Module):
def __init__(self, base_model):
super().__init__()
self.bert = base_model
self.safety_head = nn.Linear(768, 2) # safe/unsafe
self.attack_head = nn.Linear(768, 4) # jailbreak/evasion/generic/injection
def forward(self, input_ids, attention_mask):
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
return_dict=True
)
pooled = outputs.last_hidden_state[:, 0, :] # Берем [CLS] токен
return {
'safety': self.safety_head(pooled),
'attack': self.attack_head(pooled)
}
# 4. Загрузка базовой модели
base_model = BertModel.from_pretrained(
peft_config.base_model_name_or_path,
add_pooling_layer=False
).to(device)
# 5. Инициализация и загрузка LoRA
model = SafetyAttackModel(base_model).to(device)
model.bert = PeftModel.from_pretrained(model.bert, MODEL_DIR)
# 6. Проверка загрузки LoRA
print("\n=== Проверка загрузки LoRA ===")
lora_params = [name for name, _ in model.bert.named_parameters() if 'lora' in name]
if lora_params:
print(f"✅ Успешно загружено {len(lora_params)} LoRA-параметров")
print("Примеры параметров:", lora_params[:3])
else:
raise RuntimeError("❌ LoRA параметры не загрузились!")
# 7. Объединение адаптеров (опционально)
model.bert = model.bert.merge_and_unload()
model.to(device)
model.eval()
# 8. Функция для предсказаний
def predict(text):
# Токенизация (убедимся, что не передаем labels)
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['safety'], dim=1)
attack_probs = torch.softmax(outputs['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()
}
}
# 9. Пример использования
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%}")
\ No newline at end of file
checkLora.py
+ 117
− 141
View file @ 79425669
# import torch
# from transformers import BertPreTrainedModel, BertModel, BertTokenizer
# from torch import nn
# from peft import PeftModel
# # Определяем ваш кастомный класс модели, как использовался при обучении
# class MultiTaskBert(BertPreTrainedModel):
# def __init__(self, config):
# super().__init__(config)
# self.bert = BertModel(config)
# self.classifier_safety = nn.Linear(config.hidden_size, 2)
# self.classifier_attack = nn.Linear(config.hidden_size, 4)
# def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
# # Переводим тензоры на устройство
# input_ids, attention_mask, labels = map(lambda x: x.to(device) if x is not None else None, [input_ids, attention_mask, labels])
# outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
# pooled_output = outputs.last_hidden_state[:, 0, :]
# logits_safety = self.classifier_safety(pooled_output)
# logits_attack = self.classifier_attack(pooled_output)
# loss = None
# if labels is not None:
# labels_safety, labels_attack = labels[:, 0], labels[:, 1]
# loss_safety = nn.CrossEntropyLoss(weight=class_weights_task1_tensor)(logits_safety, labels_safety)
# loss_attack = nn.CrossEntropyLoss(weight=class_weights_task2_tensor)(logits_attack, labels_attack)
# loss = loss_safety + loss_attack
# return {'logits_safety': logits_safety, 'logits_attack': logits_attack, 'loss': loss}
# # Устройство
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# # Загружаем кастомную модель (без LoRA) из базового чекпоинта
# base_model = MultiTaskBert.from_pretrained("bert-base-uncased").to(device)
# # Загружаем модель с LoRA-адаптером (она была сохранена с помощью trainer.save_model)
# model = PeftModel.from_pretrained(base_model, "./fine-tuned-bert-lora_new").to(device)
# # Загружаем токенизатор (он был сохранён вместе с моделью)
# tokenizer = BertTokenizer.from_pretrained("./fine-tuned-bert-lora_new")
# # Пример инференса
# # text = "How to kill my mom."
# # inputs = tokenizer(text, return_tensors="pt")
# # inputs = {k: v.to(device) for k, v in inputs.items()}
# # with torch.no_grad():
# # outputs = model(**inputs)
# # # Получаем предсказания: для безопасности – индекс 0 (safe) или 1 (unsafe),
# # # для типа атаки – 0: jailbreak, 1: evasion, 2: generic attack, 3: other
# # pred_safety = torch.argmax(outputs['logits_safety'], dim=1).item()
# # pred_attack = torch.argmax(outputs['logits_attack'], dim=1).item()
# # safety_mapping = {0: "safe", 1: "unsafe"}
# # attack_mapping = {0: "jailbreak", 1: "evasion", 2: "generic attack", 3: "other"}
# # # Красивый вывод результатов
# # print(f"Промпт: {text}")
# # print(f"Безопасность: {safety_mapping.get(pred_safety, 'unknown')}")
# # print(f"Тип атаки: {attack_mapping.get(pred_attack, 'unknown')}")
# prompts = [
# "How to kill my mom.",
# "I love my dog.",
# "Write a story about a brave knight."
# ]
# # Проходим по каждому промпту
# for prompt in prompts:
# # Токенизируем промпт
# inputs = tokenizer(prompt, return_tensors="pt")
# inputs = {k: v.to(device) for k, v in inputs.items()}
# # Получаем вывод модели
# with torch.no_grad():
# outputs = model(**inputs)
# # Вычисляем вероятности через softmax
# safety_logits = outputs['logits_safety']
# attack_logits = outputs['logits_attack']
# safety_probs = F.softmax(safety_logits, dim=1)
# attack_probs = F.softmax(attack_logits, dim=1)
# # Находим индекс наибольшей вероятности для каждой задачи
# safety_pred = torch.argmax(safety_probs, dim=1).item()
# attack_pred = torch.argmax(attack_probs, dim=1).item()
# # Извлекаем вероятность выбранной категории (в процентах)
# safety_prob = safety_probs[0, safety_pred].item() * 100
# attack_prob = attack_probs[0, attack_pred].item() * 100
# # Вывод результатов
# print(f"Промпт: {prompt}")
# print(f"Безопасность: {safety_mapping[safety_pred]} (Вероятность: {safety_prob:.2f}%)")
# print(f"Тип атаки: {attack_mapping[attack_pred]} (Вероятность: {attack_prob:.2f}%)")
# print("-" * 50)
import torch
from transformers import BertPreTrainedModel, BertModel, BertTokenizer
from transformers import BertTokenizer, BertModel
from peft import PeftModel, PeftConfig
from torch import nn
from peft import PeftModel
# Определяем кастомную модель
class MultiTaskBert(BertPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.bert = BertModel(config)
self.classifier_safety = nn.Linear(config.hidden_size, 2)
self.classifier_attack = nn.Linear(config.hidden_size, 4)
def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
pooled_output = outputs.last_hidden_state[:, 0, :]
logits_safety = self.classifier_safety(pooled_output)
logits_attack = self.classifier_attack(pooled_output)
return {'logits_safety': logits_safety, 'logits_attack': logits_attack}
import os
# 1. Проверка устройства и файлов модели
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Загружаем базовую модель и затем оборачиваем её LoRA-адаптером
base_model = MultiTaskBert.from_pretrained("bert-base-uncased").to(device)
model = PeftModel.from_pretrained(base_model, "./fine-tuned-bert-lora_new").to(device)
tokenizer = BertTokenizer.from_pretrained("./fine-tuned-bert-lora_new")
# Пример инференса
text = "How to kill my mom."
inputs = tokenizer(text, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
pred_safety = torch.argmax(outputs['logits_safety'], dim=1).item()
pred_attack = torch.argmax(outputs['logits_attack'], dim=1).item()
safety_mapping = {0: "safe", 1: "unsafe"}
attack_mapping = {0: "jailbreak", 1: "evasion", 2: "generic attack", 3: "other"}
print(f"Промпт: {text}")
print(f"Безопасность: {safety_mapping.get(pred_safety, 'unknown')}")
print(f"Тип атаки: {attack_mapping.get(pred_attack, 'unknown')}")
\ No newline at end of file
# MODEL_DIR = "./fine-tuned-bert-lora_new"
./micro_no_cross_fine_tuned
# Проверяем наличие всех необходимых файлов
required_files = [
'adapter_config.json',
'adapter_model.safetensors',
'tokenizer_config.json',
'vocab.txt'
]
for file in required_files:
if not os.path.exists(os.path.join(MODEL_DIR, file)):
raise FileNotFoundError(f"Не найден файл {file} в {MODEL_DIR}")
# 2. Загрузка токенизатора и конфига
tokenizer = BertTokenizer.from_pretrained(MODEL_DIR)
peft_config = PeftConfig.from_pretrained(MODEL_DIR)
# 3. Определение архитектуры модели (должно совпадать с обучением)
class SafetyAttackModel(nn.Module):
def __init__(self, base_model):
super().__init__()
self.bert = base_model
self.safety_head = nn.Linear(768, 2) # safe/unsafe
self.attack_head = nn.Linear(768, 4) # jailbreak/evasion/generic/injection
def forward(self, input_ids, attention_mask):
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
return_dict=True
)
pooled = outputs.last_hidden_state[:, 0, :] # Берем [CLS] токен
return {
'safety': self.safety_head(pooled),
'attack': self.attack_head(pooled)
}
# 4. Загрузка базовой модели
base_model = BertModel.from_pretrained(
peft_config.base_model_name_or_path,
add_pooling_layer=False
).to(device)
# 5. Инициализация и загрузка LoRA
model = SafetyAttackModel(base_model).to(device)
model.bert = PeftModel.from_pretrained(model.bert, MODEL_DIR)
# 6. Проверка загрузки LoRA
print("\n=== Проверка загрузки LoRA ===")
lora_params = [name for name, _ in model.bert.named_parameters() if 'lora' in name]
if lora_params:
print(f"✅ Успешно загружено {len(lora_params)} LoRA-параметров")
print("Примеры параметров:", lora_params[:3])
else:
raise RuntimeError("❌ LoRA параметры не загрузились!")
# 7. Объединение адаптеров (опционально)
model.bert = model.bert.merge_and_unload()
model.to(device)
model.eval()
# 8. Функция для предсказаний
def predict(text):
# Токенизация (убедимся, что не передаем labels)
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['safety'], dim=1)
attack_probs = torch.softmax(outputs['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()
}
}
# 9. Пример использования
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%}")
\ No newline at end of file
microzapusk.py
+ 1
− 1
View file @ 79425669
......@@ -140,7 +140,7 @@ for fold, (train_idx, val_idx) in enumerate(kf.split(data)):
# Сохранение новой лучшей модели
model.save_pretrained(os.path.join(OUTPUT_DIR, 'best_model'))
# model.save_adapter(os.path.join(OUTPUT_DIR, 'best_model'), "lora_adapters")
model.save_adapter(os.path.join(OUTPUT_DIR, 'best_model'), "lora_adapters")
tokenizer.save_pretrained(os.path.join(OUTPUT_DIR, 'best_model'))
# Итоговые результаты
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment