Commit f41e9854 authored by Мазур Грета Евгеньевна's avatar Мазур Грета Евгеньевна
Browse files

micro zapusk no cross

parent 098d5319
No related merge requests found
Showing with 168 additions and 4 deletions
+168 -4
# 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 import torch
from transformers import BertTokenizer, BertForSequenceClassification from transformers import BertTokenizer, BertForSequenceClassification
from peft import get_peft_model, LoraConfig, TaskType from peft import get_peft_model, LoraConfig, TaskType
...@@ -37,8 +113,14 @@ def predict(text): ...@@ -37,8 +113,14 @@ def predict(text):
# Ожидаем, что модель вернет логиты # Ожидаем, что модель вернет логиты
with torch.no_grad(): with torch.no_grad():
outputs = model(**inputs) 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 для получения вероятностей
safety_probs = torch.softmax(logits_safety, dim=1)
attack_probs = torch.softmax(logits_attack, dim=1)
return { return {
'safety': { 'safety': {
......
# 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 import torch
from transformers import BertTokenizer, BertForSequenceClassification from transformers import BertTokenizer, BertForSequenceClassification
from peft import get_peft_model, LoraConfig, TaskType from peft import get_peft_model, LoraConfig, TaskType
...@@ -37,8 +113,14 @@ def predict(text): ...@@ -37,8 +113,14 @@ def predict(text):
# Ожидаем, что модель вернет логиты # Ожидаем, что модель вернет логиты
with torch.no_grad(): with torch.no_grad():
outputs = model(**inputs) 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 для получения вероятностей
safety_probs = torch.softmax(logits_safety, dim=1)
attack_probs = torch.softmax(logits_attack, dim=1)
return { return {
'safety': { 'safety': {
......
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