# 2 модуль - Генератор парафраза на основе llama.cpp

from openai import OpenAI

class ParaphraseGenerator:
    """
    Класс для использования генеративной модели на базе llama.cpp

    Использование: ParaphraseGenerator([температура], [максимальное количество токенов])
    """

    def __init__(self, temperature=0.8, max_tokens=512):
        self.client = OpenAI(base_url='http://127.0.0.1:8080', api_key='')
        self.temperature = temperature
        self.max_tokens = max_tokens

    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)