diff --git a/.ipynb_checkpoints/ULTRAMegaOB-checkpoint.py b/.ipynb_checkpoints/ULTRAMegaOB-checkpoint.py
index 103a15bcbc48b1a7f7922aeb801ac1d5c52a0db8..66ae3915316c0681a46181260046025c861aa944 100644
--- a/.ipynb_checkpoints/ULTRAMegaOB-checkpoint.py
+++ b/.ipynb_checkpoints/ULTRAMegaOB-checkpoint.py
@@ -180,13 +180,12 @@ def compute_metrics(p):
 
 def augment_text(text, num_augments):
     """Генерация аугментированных примеров с проверками"""
-    
-    if len(text) > 1000:  # Слишком длинные тексты плохо аугментируются
+    if len(text) > 1000:
         logger.debug(f"Текст слишком длинный для аугментации: {len(text)} символов")
         return [text]
     
     if not isinstance(text, str) or len(text.strip()) < 10:
-        return []
+        return [text]
         
     text = text.replace('\n', ' ').strip()
     
@@ -208,12 +207,9 @@ def augment_text(text, num_augments):
         # Обратный перевод
         if len(augmented) < num_augments:
             try:
-                # Определяем язык текста
                 if any(cyr_char in text for cyr_char in 'абвгдеёжзийклмнопрстуфхцчшщъыьэюя'):
-                    # Для русских текстов
                     tr_augs = translation_aug_ru.augment(text, n=num_augments-len(augmented))
                 else:
-                    # Для английских/других текстов
                     tr_augs = translation_aug.augment(text, n=num_augments-len(augmented))
                     
                 if tr_augs:
@@ -227,7 +223,7 @@ def augment_text(text, num_augments):
             logger.debug(f"Не удалось аугментировать текст: {text[:50]}...")
             return [text]
             
-        augmented = list(set(augmented))  # Удаление дубликатов
+        augmented = list(set(augmented))
         return list(augmented)[:num_augments] if augmented else [text]
     except Exception as e:
         logger.error(f"Критическая ошибка аугментации: {str(e)}")
@@ -259,8 +255,8 @@ def balance_attack_types(unsafe_data):
             
             augmented = subset.sample(n=num_augments, replace=True)
             augmented['prompt'] = augmented['prompt'].apply(
-                lambda x: (augs := augment_text(x, 1)) and augs[0] if augs else x
-            )
+            lambda x: augment_text(x, 1)[0]  # Просто берем первый элемент возвращаемого списка
+                )
             
             # Логирование аугментированных примеров
             logger.info(f"\nАугментация для {attack_type}:")
diff --git a/ULTRAMegaOB.py b/ULTRAMegaOB.py
index 103a15bcbc48b1a7f7922aeb801ac1d5c52a0db8..66ae3915316c0681a46181260046025c861aa944 100644
--- a/ULTRAMegaOB.py
+++ b/ULTRAMegaOB.py
@@ -180,13 +180,12 @@ def compute_metrics(p):
 
 def augment_text(text, num_augments):
     """Генерация аугментированных примеров с проверками"""
-    
-    if len(text) > 1000:  # Слишком длинные тексты плохо аугментируются
+    if len(text) > 1000:
         logger.debug(f"Текст слишком длинный для аугментации: {len(text)} символов")
         return [text]
     
     if not isinstance(text, str) or len(text.strip()) < 10:
-        return []
+        return [text]
         
     text = text.replace('\n', ' ').strip()
     
@@ -208,12 +207,9 @@ def augment_text(text, num_augments):
         # Обратный перевод
         if len(augmented) < num_augments:
             try:
-                # Определяем язык текста
                 if any(cyr_char in text for cyr_char in 'абвгдеёжзийклмнопрстуфхцчшщъыьэюя'):
-                    # Для русских текстов
                     tr_augs = translation_aug_ru.augment(text, n=num_augments-len(augmented))
                 else:
-                    # Для английских/других текстов
                     tr_augs = translation_aug.augment(text, n=num_augments-len(augmented))
                     
                 if tr_augs:
@@ -227,7 +223,7 @@ def augment_text(text, num_augments):
             logger.debug(f"Не удалось аугментировать текст: {text[:50]}...")
             return [text]
             
-        augmented = list(set(augmented))  # Удаление дубликатов
+        augmented = list(set(augmented))
         return list(augmented)[:num_augments] if augmented else [text]
     except Exception as e:
         logger.error(f"Критическая ошибка аугментации: {str(e)}")
@@ -259,8 +255,8 @@ def balance_attack_types(unsafe_data):
             
             augmented = subset.sample(n=num_augments, replace=True)
             augmented['prompt'] = augmented['prompt'].apply(
-                lambda x: (augs := augment_text(x, 1)) and augs[0] if augs else x
-            )
+            lambda x: augment_text(x, 1)[0]  # Просто берем первый элемент возвращаемого списка
+                )
             
             # Логирование аугментированных примеров
             logger.info(f"\nАугментация для {attack_type}:")