add dataset

.gitignore

@@ -1,5 +1,5 @@
 .venv/
 llama.cpp/
-NEREL/
+rupaws/
 __pycache__/
 models/

main.py

@@ -1,5 +1,6 @@
 # Точка входа программы
 
+import csv
 from modules.NER import NER
 from modules.paraGenerator import ParaphraseGenerator
 from modules.validator import validator
@@ -7,14 +8,22 @@ from modules.validator import validator
 ner = NER()
 pg = ParaphraseGenerator()
 
-srcText = 'Добрый день, я, Сидоров Иван Иванович. Прошу перевести сто тысяч рублей Якову Петру Игнатьевичу в Москву.'
+datasetPath = 'rupaws/rupaws_wiki_test.csv'
+
+srcTexts = []
+with open(datasetPath, 'r') as f:
+    reader = csv.reader(f, delimiter=';')
+    for row in reader:
+        srcTexts.append(row[1])
 
 def main(srcText):
     srcEntities = ner.extract_entities(srcText) # поиск сущностей
+    print(f'ИСХОДНЫЕ СУЩНОСТИ ~> {srcEntities}')
     paraphrase = pg.generate(srcText, srcEntities) # генерация парафраза
     paraEntities = ner.extract_entities(paraphrase) # поиск сущностей в парафразе
     return validator(srcText, paraphrase, srcEntities, paraEntities) # валидация
 
-result = main(srcText)
-print(f'ИСХОДНЫЙ ТЕКСТ: {srcText}')
-print(f'СГЕНЕРИРОВАННЫЙ ТЕКСТ: {result}')
+for i in range(10): # проверим 10 текстов
+    text = srcTexts[i]
+    print(f'ИСХОДНЫЙ ТЕКСТ ~> {text}')
+    print(f'ПАРАФРАЗИРОВАННЫЙ ТЕКСТ ~> {main(text)}\n')

modules/paraGenerator.py

@@ -39,6 +39,7 @@ class ParaphraseGenerator:
             f'Перефразируй следующий текст, сохранив все именованные сущности "{', '.join(entity for entity in entities)}" в точности такими же, как в оригинале. '
              'Не изменяй эти фрагменты текста, не заменяй их синонимами, не переставляй слова внутри них. '
              'Можешь изменять грамматическую структуру предложения, порядок слов, использовать синонимы для'
-            f'остальных частей текста, но именованные сущности должны остаться неизменными. Исходный текст: "{srcText}"'
+            f'остальных частей текста, но именованные сущности должны остаться неизменными.'
+            f'Исходный текст: "{srcText}". Выводи только сгенерированный текст.'
         )
         return self.generateByPrompt(prompt)

modules/validator.py

@@ -56,14 +56,13 @@ def validator(srcText, paraphrase, srcEntities, paraEntities):
     """
     ce = compare_entities(srcEntities, paraEntities)
     if ce[0]:
+        print(f'ПАРАФРАЗИРОВАННЫЕ СУЩНОСТИ ~> {paraEntities}')
         return paraphrase # если всё нормально, возвращаем текст в неизменном виде
     
-    print(f'Произошла потеря сущностей!')
-    print(f'Исходные сущности: {srcEntities}')
-    print(f'Сгенерируемые сущности: {paraEntities}')
-    print(f'Потеря сущностей: {ce[1]}')
     # даём 5 попыток на восстановление
     for i in range(5):
+        print(f'ПОТЕРЯННЫЕ СУЩНОСТИ ~> {ce[1]}')
+        print(f'ПОПЫТКА ВОССТАНОВЛЕНИЯ ~> {i+1}')
         regen_prompt = (
         f'При перефразировании текста "{srcText}" из списка элементов "{', '.join(entity for entity in srcEntities)}"'
         f'были утеряны или изменены следующие важные элементы: "{', '.join(e for e in ce[1])}". '
@@ -72,14 +71,7 @@ def validator(srcText, paraphrase, srcEntities, paraEntities):
         newParaphrase = pg.generateByPrompt(regen_prompt)
         paraEntities = ner.extract_entities(newParaphrase)
         ce = compare_entities(srcEntities, paraEntities)
-        print(f'Попытка восстановления: {i+1}')
-        print(f'Исходные сущности: {srcEntities}')
-        print(f'Сгенерируемые сущности: {paraEntities}')
-        print(f'Потеря сущностей: {ce[1]}')
         if (ce[0]): # если сравнение дало True, выходим из цикла
+            print(f'ПАРАФРАЗИРОВАННЫЕ СУЩНОСТИ ~> {paraEntities}')
             return newParaphrase
     return None
-
-# a = set(['a','b', 'c', 'd'])
-# b = set(['a','b'])
-# validator('123', '123', a, b)