pereobuch2

.ipynb_checkpoints/superPereObuch-checkpoint.py

@@ -36,26 +36,31 @@ def load_and_balance_data():
     # Балансировка для редких классов атак
     attack_types = unsafe_data['type'].value_counts()
     
-    # Увеличиваем количество редких классов
+    # Увеличиваем количество редких классов с заменой (replace=True)
     balanced_unsafe = pd.concat([
         unsafe_data[unsafe_data['type'] == 'evasion'].sample(
-            n=int(len(unsafe_data)*0.1),  # 10% от общего числа unsafe
-            replace=True,
+            n=max(1, int(len(unsafe_data)*0.1)),  # Гарантируем хотя бы 1 пример
+            replace=True,  # Разрешаем повторения
             random_state=42
         ),
         unsafe_data[unsafe_data['type'] == 'generic attack'].sample(
-            n=int(len(unsafe_data)*0.05),  # 5% от общего числа unsafe
-            replace=True,
+            n=max(1, int(len(unsafe_data)*0.05)),  # Гарантируем хотя бы 1 пример
+            replace=True,  # Разрешаем повторения
             random_state=42
         ),
         unsafe_data[unsafe_data['type'].isin(['jailbreak', 'injection'])]
     ])
     
-    # Финалный датасет (50/50 safe/unsafe)
-    balanced_data = pd.concat([
-        safe_data.sample(n=len(balanced_unsafe)), 
-        balanced_unsafe
-    ]).sample(frac=1, random_state=42)
+    # Берем выборку безопасных примеров с заменой, если нужно
+    n_samples = min(len(safe_data), len(balanced_unsafe))
+    balanced_safe = safe_data.sample(
+        n=n_samples,
+        replace=len(safe_data) < len(balanced_unsafe),  # Разрешаем замену только если нужно
+        random_state=42
+    )
+    
+    # Финалный датасет
+    balanced_data = pd.concat([balanced_safe, balanced_unsafe]).sample(frac=1, random_state=42)
     
     print("\nРаспределение после балансировки:")
     print("Безопасность:", balanced_data['safety'].value_counts(normalize=True))
@@ -63,6 +68,7 @@ def load_and_balance_data():
     
     return balanced_data
 
+    
 # 2. Токенизация с правильными именами колонок
 def tokenize_data(tokenizer, df):
     df = df.dropna(subset=['prompt'])
@@ -183,10 +189,26 @@ def compute_metrics(p):
 
 def main():
     # 1. Подготовка данных
+
     data = load_and_balance_data()
+    
+    # Проверка что данные не пустые
+    if len(data) == 0:
+        raise ValueError("После балансировки получился пустой датасет. Проверьте исходные данные.")
+    
+    # Проверка распределения классов
+    print("\nПроверка распределения перед обучением:")
+    print("Safe:", len(data[data['safety'] == 'safe']))
+    print("Unsafe:", len(data[data['safety'] == 'unsafe']))
+    
     train_data, test_data = train_test_split(data, test_size=0.2, stratify=data['safety'])
     train_data, val_data = train_test_split(train_data, test_size=0.1, stratify=train_data['safety'])
     
+    # # ... остальной код
+    # data = load_and_balance_data()
+    # train_data, test_data = train_test_split(data, test_size=0.2, stratify=data['safety'])
+    # train_data, val_data = train_test_split(train_data, test_size=0.1, stratify=train_data['safety'])
+    
     # 2. Токенизация
     tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
     train_dataset = tokenize_data(tokenizer, train_data)

superPereObuch.py

@@ -36,26 +36,31 @@ def load_and_balance_data():
     # Балансировка для редких классов атак
     attack_types = unsafe_data['type'].value_counts()
     
-    # Увеличиваем количество редких классов
+    # Увеличиваем количество редких классов с заменой (replace=True)
     balanced_unsafe = pd.concat([
         unsafe_data[unsafe_data['type'] == 'evasion'].sample(
-            n=int(len(unsafe_data)*0.1),  # 10% от общего числа unsafe
-            replace=True,
+            n=max(1, int(len(unsafe_data)*0.1)),  # Гарантируем хотя бы 1 пример
+            replace=True,  # Разрешаем повторения
             random_state=42
         ),
         unsafe_data[unsafe_data['type'] == 'generic attack'].sample(
-            n=int(len(unsafe_data)*0.05),  # 5% от общего числа unsafe
-            replace=True,
+            n=max(1, int(len(unsafe_data)*0.05)),  # Гарантируем хотя бы 1 пример
+            replace=True,  # Разрешаем повторения
             random_state=42
         ),
         unsafe_data[unsafe_data['type'].isin(['jailbreak', 'injection'])]
     ])
     
-    # Финалный датасет (50/50 safe/unsafe)
-    balanced_data = pd.concat([
-        safe_data.sample(n=len(balanced_unsafe)), 
-        balanced_unsafe
-    ]).sample(frac=1, random_state=42)
+    # Берем выборку безопасных примеров с заменой, если нужно
+    n_samples = min(len(safe_data), len(balanced_unsafe))
+    balanced_safe = safe_data.sample(
+        n=n_samples,
+        replace=len(safe_data) < len(balanced_unsafe),  # Разрешаем замену только если нужно
+        random_state=42
+    )
+    
+    # Финалный датасет
+    balanced_data = pd.concat([balanced_safe, balanced_unsafe]).sample(frac=1, random_state=42)
     
     print("\nРаспределение после балансировки:")
     print("Безопасность:", balanced_data['safety'].value_counts(normalize=True))
@@ -63,6 +68,7 @@ def load_and_balance_data():
     
     return balanced_data
 
+    
 # 2. Токенизация с правильными именами колонок
 def tokenize_data(tokenizer, df):
     df = df.dropna(subset=['prompt'])
@@ -183,10 +189,26 @@ def compute_metrics(p):
 
 def main():
     # 1. Подготовка данных
+
     data = load_and_balance_data()
+    
+    # Проверка что данные не пустые
+    if len(data) == 0:
+        raise ValueError("После балансировки получился пустой датасет. Проверьте исходные данные.")
+    
+    # Проверка распределения классов
+    print("\nПроверка распределения перед обучением:")
+    print("Safe:", len(data[data['safety'] == 'safe']))
+    print("Unsafe:", len(data[data['safety'] == 'unsafe']))
+    
     train_data, test_data = train_test_split(data, test_size=0.2, stratify=data['safety'])
     train_data, val_data = train_test_split(train_data, test_size=0.1, stratify=train_data['safety'])
     
+    # # ... остальной код
+    # data = load_and_balance_data()
+    # train_data, test_data = train_test_split(data, test_size=0.2, stratify=data['safety'])
+    # train_data, val_data = train_test_split(train_data, test_size=0.1, stratify=train_data['safety'])
+    
     # 2. Токенизация
     tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
     train_dataset = tokenize_data(tokenizer, train_data)