add validator

NER.py

@@ -1,27 +0,0 @@
-# 1 модуль - Распознавание именованных сущностей (NER)
-
-import spacy
-
-class NER:
-    """
-    Класс для выделения именованных сущностей из текста с помощью библиотеки spaCy.
-    """
-    def __init__(self):
-        self.nlp = spacy.load("ru_core_news_lg")
-    
-    def extract_entities(self, text):
-        """
-        Выделение именованных сущностей из текста
-        Использование: text (<текст>)
-        Возвращает: List[Dict[text, type]] - список словарей, каждый из которых содержит информацию об одной сущности
-        """
-        doc = self.nlp(text)
-        entities = []
-        for ent in doc.ents:
-            entities.append({
-                'text': ent.text,        # Текст сущности
-                'type': ent.label_,      # Тип сущности (PER, LOC, ORG, DATE и т.д.)
-                'start': ent.start_char, # Начальная позиция в тексте
-                'end': ent.end_char      # Конечная позиция в тексте
-            })
-        return entities

README.md

@@ -1,11 +1,15 @@
 ## Программная реализация ВКР "Алгоритм парафразирования для улучшения извлечения сущностей с использованием больших языковых моделей"
 ### Установка библиотек
-#### для NER
+#### Для NER и validator
 ```
 pip install spacy
 python -m spacy download ru_core_news_lg
 ```
-#### для paraGenerator
+#### Для paraGenerator
 ```
 CMAKE_ARGS="-DGGML_CUDA=on -DLLAMA_OPENSSL=ON" pip install llama-cpp-python
 ```
+#### Для validator
+```
+pip install pymorphy3
+```

main.py

@@ -1,26 +1,31 @@
 # Точка входа программы
 
-from NER import NER
+from modules.NER import NER
-from paraGenerator import ParaphraseGenerator
+from modules.paraGenerator import ParaphraseGenerator
+#from modules.validator import validator
 
-## Запуск рабочих инструментов
 ner = NER()
 pg = ParaphraseGenerator()
-###
 
 srcText = 'Добрый день, я, Сидоров Иван Иванович. Прошу перевести сто тысяч рублей Якову Петру Игнатьевичу в Москву.'
 
-entities1 = ner.extract_entities(srcText)
+def main(srcText):
-print(entities1)
+    # поиск сущностей
+    srcEntities = ner.extract_entities(srcText)
+    print(srcEntities)
 
-#paraphrase1 = pg.generate(srcText, entities1)
+    # генерация парафраза
-#print(paraphrase1)
+    paraphrase = pg.generate(srcText, srcEntities)
+    print(paraphrase)
+
+    # поиск сущностей в парафразе
+    paraEntities = ner.extract_entities(paraphrase)
+    print(paraEntities)
+
+    # Валидация
+ #   return validator(srcText, paraphrase, srcEntities, paraEntities)
+
+result = main(srcText)
+print(result)
 
-#entities2 = ner.extract_entities(paraphrase1)
-#print(entities2)
 
-gg = ()
-gg.append(1)
-gg.append(2)
-gg.append(1)
-print(gg)

modules/NER.py

@@ -0,0 +1,24 @@
+# 1 модуль - Распознавание именованных сущностей (NER)
+
+import spacy
+
+class NER:
+    """
+    Класс для выделения именованных сущностей из текста с помощью библиотеки spaCy.
+    """
+    def __init__(self):
+        self.nlp = spacy.load('ru_core_news_lg')
+    
+    def extract_entities(self, text):
+        """
+        Выделение именованных сущностей из текста.
+
+        Использование: text (<текст>)
+
+        Возвращает set сущностей без повторения
+        """
+        doc = self.nlp(text)
+        entities = set()
+        for ent in doc.ents:
+            entities.add(ent.text)
+        return entities

modules/paraGenerator.py

@@ -0,0 +1,42 @@
+# 2 модуль - Генератор парафраза на основе llama.cpp
+
+from openai import OpenAI
+
+class ParaphraseGenerator:
+    """
+    Класс для использования генеративной модели на базе llama.cpp
+
+    Использование: ParaphraseGenerator([температура], [максимальное количество токенов])
+    """
+
+    def __init__(self, temperature=0.8, max_tokens=200):
+        self.client = OpenAI(base_url='http://127.0.0.1:8080', api_key='')
+
+    def generateByPrompt(self, prompt):
+        try:
+            response = self.client.chat.completions.create(
+                model = '',
+                messages = [{'role': 'user', 'content': prompt}],
+                temperature = self.temperature,
+                max_tokens = self.max_tokens
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            print(f'Ошибка генерации: {e}')
+            return None
+
+    def generate(self, srcText, entities):
+        """
+        Генерация парафраза.
+
+        Использование: generate(<исходный текст>, <сущности>)
+        
+        Возвращает set сущностей (без повторения)
+        """
+        prompt = (
+            f'Перефразируй следующий текст, сохранив все именованные сущности "{', '.join(entity for entity in entities)}" в точности такими же, как в оригинале. '
+             'Не изменяй эти фрагменты текста, не заменяй их синонимами, не переставляй слова внутри них. '
+             'Можешь изменять грамматическую структуру предложения, порядок слов, использовать синонимы для'
+            f'остальных частей текста, но именованные сущности должны остаться неизменными. Исходный текст: "{srcText}"'
+        )
+        return self.generateByPrompt(prompt)

modules/validator.py

@@ -0,0 +1,74 @@
+# 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)

paraGenerator.py

@@ -1,33 +0,0 @@
-# 2 модуль - Генератор парафраза на основе llama.cpp
-
-from openai import OpenAI
-
-class ParaphraseGenerator:
-    """
-    Класс для использования генеративной модели на базе llama.cpp
-    """
-    def __init__(self):
-        self.client = OpenAI(base_url="http://127.0.0.1:8080", api_key="")
-
-    def generate(self, srcText, entities):
-        """
-        Генерация парафраза по промпту.
-        Использование: generate(<исходный текст>, <сущности>)
-        """
-        entityList = ', '.join(entity['text'] for entity in entities) # генерация списка выделенных сущностей
-        prompt = 'Перефразируй следующий текст, сохранив все именованные сущности "' + entityList + '" в точности такими же, как в оригинале. ' \
-                 'Не изменяй эти фрагменты текста, не заменяй их синонимами, не переставляй слова внутри них. ' \
-                 'Можешь изменять грамматическую структуру предложения, порядок слов, использовать синонимы для' \
-                 'остальных частей текста, но именованные сущности должны остаться неизменными. Исходный текст: "' + srcText + '"'
-        try:
-            response = self.client.chat.completions.create(
-                model = "",
-                messages = [{"role": "user", "content": prompt}],
-                temperature = 0.8,
-                top_p = 0.95,
-                max_tokens = 200
-            )
-            return response.choices[0].message.content
-        except Exception as e:
-            print(f"Ошибка генерации: {e}")
-            return None