diff --git a/modules/validator.py b/modules/validator.py index 7388790..b53545d 100644 --- a/modules/validator.py +++ b/modules/validator.py @@ -12,6 +12,8 @@ def compare_entities(original, generated): """ Сравнивает два списка сущностей. Новые сущности допускаются. + Использование: compare_entities(<сущности исходного текста>, <сущности сгенерированного текста>) + Возвращает: - True - всё на месте - False - что то потерялось @@ -56,16 +58,17 @@ def validator(srcText, srcEntities, paraEntities): if ce[0]: return srcText # если всё нормально, возвращаем текст в неизменном виде - regen_prompt = ( + # даём 3 попытки на восстановление + for _ in range(3): + regen_prompt = ( f'При перефразировании текста "{srcText}" из списка элементов "{', '.join(entity for entity in srcEntities)}"' f'были утеряны или изменены следующие важные элементы: "{', '.join(e for e in ce[1])}". ' 'Перефразируй исходный текст заново, обратив особое внимание на сохранение этих элементов. Выведи только текст.' - ) - - for _ in range(3): + ) newParaphrase = pg.generateByPrompt(regen_prompt) paraEntities = ner.extract_entities(newParaphrase) - if (compare_entities(srcEntities, paraEntities)): + ce = compare_entities(srcEntities, paraEntities) + if (ce[0]): # если сравнение дало True, выходим из цикла return newParaphrase return None