75 lines
3.2 KiB
Python
75 lines
3.2 KiB
Python
# 3 модуль - Валидация и восстановление сущностей
|
||
|
||
from NER import NER
|
||
from paraGenerator import ParaphraseGenerator
|
||
from pymorphy3 import MorphAnalyzer
|
||
|
||
ner = NER()
|
||
pg = ParaphraseGenerator()
|
||
morph = MorphAnalyzer()
|
||
|
||
def compare_entities(original, generated):
|
||
"""
|
||
Сравнивает два списка сущностей. Новые сущности допускаются.
|
||
|
||
Возвращает:
|
||
- True - всё на месте
|
||
- False - что то потерялось
|
||
"""
|
||
def normalize(text):
|
||
"""
|
||
Приводит слово к нормальной форме. Пример: Ивану --> Иван
|
||
|
||
Возвращает:
|
||
- Совпадают или нет сущности (bool)
|
||
- Какие сущности потерялись (set)
|
||
"""
|
||
return morph.parse(text)[0].normal_form
|
||
|
||
if original.issubset(generated): # если original является подмножеством множества generated
|
||
return True, set()
|
||
|
||
# если нет, проверяем дополнительно, вдруг кейс по типу "Иван" - "Ивану"
|
||
orig_norm = {normalize(e) for e in original} # нормализуем списки
|
||
gen_norm = {normalize(e) for e in generated}
|
||
|
||
if orig_norm.issubset(gen_norm):
|
||
return True, set()
|
||
|
||
# если по прежнему false, то ищем потерянные сущности
|
||
lost = set()
|
||
for o in original:
|
||
if (normalize(o) not in gen_norm):
|
||
lost.add(o)
|
||
return False, lost
|
||
|
||
def validator(srcText, srcEntities, paraEntities):
|
||
"""
|
||
Использование: validator(<исходный текст>, <сущности исходного текста>, <сущности перефразированного текста>)
|
||
|
||
Возвращает:
|
||
- Исходный текст если сущности сохранены
|
||
- Изменённый текст, если сущности не сохранены и были восстановлены
|
||
- None, если сущности не удалось восстановить с трёх раз
|
||
"""
|
||
ce = compare_entities(srcEntities, paraEntities)
|
||
if ce[0]:
|
||
return srcText # если всё нормально, возвращаем текст в неизменном виде
|
||
|
||
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)):
|
||
return newParaphrase
|
||
return None
|
||
|
||
# a = set(['a','b', 'c', 'd'])
|
||
# b = set(['a','b'])
|
||
# validator('123', '123', a, b)
|