# Установка необходимых библиотек
# Убедитесь, что вы установили все необходимые библиотеки перед запуском кода.
# Вы можете установить их с помощью pip:
# pip install tensorflow numpy nltk

import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, LSTM, Dense, Embedding, Concatenate, Layer
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
import numpy as np
import nltk
from nltk.corpus import stopwords
import string

# Загрузка ресурсов NLTK
nltk.download('stopwords')

# 1. Подготовка исходных данных
questions = [
    "Как тебя зовут?",
    "Сколько тебе лет?",
    "Откуда ты?",
    "Чем ты занимаешься?",
    "Какой сегодня день?",
    "Какая сейчас погода?",
    "Что ты умеешь?",
    "Кто твой создатель?",
    "Как мне добраться до ближайшего магазина?",
    "Можешь рассказать анекдот?",
    "Какое сегодня число?",
    "Какие у тебя хобби?",
    "Что ты знаешь о машинном обучении?",
    "Можешь помочь с домашним заданием?",
    "Какие языки программирования ты знаешь?",
    "Какова столица Франции?",
    "Расскажи мне о космосе.",
    "Что такое искусственный интеллект?",
    "Можешь порекомендовать книгу?",
    "Как приготовить борщ?",
    "Что ты знаешь о Python?",
    "Какую музыку ты предпочитаешь?",
    "Есть ли у тебя любимый фильм?",
    "Можешь рассказать о своей работе?",
    "Каковы твои планы на будущее?"
]

answers = [
    "Меня зовут Алекс.",
    "Мне 25 лет.",
    "Я из Москвы.",
    "Я изучаю машинное обучение.",
    "Сегодня среда.",
    "Сегодня солнечно с небольшой облачностью.",
    "Я могу вести беседу и отвечать на вопросы.",
    "Меня создала команда разработчиков.",
    "Чтобы добраться до магазина, вам нужно идти прямо по улице и повернуть направо на втором светофоре.",
    "Конечно! Почему программисты любят природу? Потому что там много багов.",
    "Сегодня 18 ноября 2024 года.",
    "Я люблю читать и программировать.",
    "Машинное обучение — это область искусственного интеллекта, которая занимается разработкой алгоритмов, позволяющих компьютерам учиться на данных.",
    "Конечно, в чем нужна помощь?",
    "Я знаком с Python, Java, C++ и другими языками.",
    "Столицей Франции является Париж.",
    "Космос — это пространство за пределами земной атмосферы, где находятся звезды, планеты и другие небесные тела.",
    "Искусственный интеллект — это способность машин выполнять задачи, требующие человеческого интеллекта.",
    "Я рекомендую книгу 'Машинное обучение для начинающих'.",
    "Для приготовления борща вам понадобятся свекла, капуста, картофель, морковь, лук и другие ингредиенты. Сначала обжарьте овощи, затем добавьте их в кипящий бульон и варите до готовности.",
    "Python — это высокоуровневый язык программирования, известный своей простотой и широким применением.",
    "Мне нравится классическая музыка и современные электронные жанры.",
    "У меня нет любимого фильма, но я могу порекомендовать известные картины.",
    "Я создан для того, чтобы помогать людям отвечать на вопросы и вести беседы.",
    "Я стремлюсь улучшаться и обучаться новым навыкам, чтобы быть полезнее для пользователей."
]

# 2. Функция препроцессинга текста
def preprocess_text(texts):
    processed_texts = []
    stop_words = set(stopwords.words('russian'))
    translator = str.maketrans('', '', string.punctuation)
    for text in texts:
        text = text.lower()  # Приведение к нижнему регистру
        text = text.translate(translator)  # Удаление пунктуации
        tokens = text.split()
        tokens = [word for word in tokens if word not in stop_words]  # Удаление стоп-слов
        processed_text = ' '.join(tokens)
        processed_texts.append(processed_text)
    return processed_texts

clean_questions = preprocess_text(questions)
clean_answers = preprocess_text(answers)

print("Препроцессинг завершен.")
print("Вопросы после препроцессинга:", clean_questions)
print("Ответы после препроцессинга:", clean_answers)

# 3. Добавление токенов <START> и <END> к ответам и создание токенизаторов
# Добавляем токены к ответам
clean_answers = ['<START> ' + answer + ' <END>' for answer in clean_answers]

# Токенизатор для вопросов
question_tokenizer = Tokenizer(oov_token='<OOV>')
question_tokenizer.fit_on_texts(clean_questions)
question_vocab_size = len(question_tokenizer.word_index) + 1  # +1 для токена OOV

# Токенизатор для ответов
answer_tokenizer = Tokenizer(oov_token='<OOV>')
answer_tokenizer.fit_on_texts(clean_answers)
answer_vocab_size = len(answer_tokenizer.word_index) + 1  # +1 для токена OOV

print(f"Размер словаря вопросов: {question_vocab_size}")
print(f"Размер словаря ответов: {answer_vocab_size}")

# 4. Преобразование текстов в последовательности чисел
question_sequences = question_tokenizer.texts_to_sequences(clean_questions)
answer_sequences = answer_tokenizer.texts_to_sequences(clean_answers)

# Определение максимальной длины последовательности
max_question_length = max(len(seq) for seq in question_sequences)
max_answer_length = max(len(seq) for seq in answer_sequences)

# Дополнение последовательностей до максимальной длины
encoder_input_data = pad_sequences(question_sequences, maxlen=max_question_length, padding='post')
decoder_input_data = pad_sequences([seq[:-1] for seq in answer_sequences], maxlen=max_answer_length - 1, padding='post')
decoder_target_data = pad_sequences([seq[1:] for seq in answer_sequences], maxlen=max_answer_length - 1, padding='post')

print("Данные подготовлены.")

# 3.a. **Определение Класса BahdanauAttention**
class BahdanauAttention(Layer):
    def __init__(self, units, **kwargs):
        super(BahdanauAttention, self).__init__(**kwargs)
        self.W1 = Dense(units)
        self.W2 = Dense(units)
        self.V = Dense(1)

    def call(self, query, values, mask=None):
        # query: decoder LSTM output (batch_size, Tq, hidden_size)
        # values: encoder LSTM output (batch_size, Tv, hidden_size)

        # Расширяем размерности для сложения
        hidden_with_time_axis = tf.expand_dims(query, 2)  # (batch_size, Tq, 1, hidden_size)

        # Вычисляем score
        score = self.V(tf.nn.tanh(self.W1(values) + self.W2(hidden_with_time_axis)))  # (batch_size, Tq, Tv, 1)

        # Применяем маску, если она предоставлена
        if mask is not None:
            # Предполагается, что mask имеет форму (batch_size, Tv)
            mask = tf.expand_dims(mask, 1)  # (batch_size, 1, Tv)
            score += (1.0 - tf.cast(mask, tf.float32)) * -1e9  # Маскируем запрещённые позиции

        # Вычисляем вес внимания
        attention_weights = tf.nn.softmax(score, axis=2)  # (batch_size, Tq, Tv, 1)

        # Вычисляем контекстный вектор
        context_vector = attention_weights * tf.expand_dims(values, 1)  # Broadcasting
        context_vector = tf.reduce_sum(context_vector, axis=2)  # (batch_size, Tq, hidden_size)

        return context_vector

# 5. Построение модели Seq2Seq с механизмом внимания
embedding_dim = 128
lstm_units = 256

# Энкодер
encoder_inputs = Input(shape=(max_question_length,), name='encoder_inputs')
encoder_embedding = Embedding(input_dim=question_vocab_size, output_dim=embedding_dim, mask_zero=True, name='encoder_embedding')(encoder_inputs)
encoder_outputs, state_h, state_c = LSTM(lstm_units, return_sequences=True, return_state=True, name='encoder_lstm')(encoder_embedding)

# Декодер
decoder_inputs = Input(shape=(max_answer_length - 1,), name='decoder_inputs')
decoder_embedding = Embedding(input_dim=answer_vocab_size, output_dim=embedding_dim, mask_zero=True, name='decoder_embedding')(decoder_inputs)
decoder_lstm = LSTM(lstm_units, return_sequences=True, return_state=True, name='decoder_lstm')
decoder_outputs, _, _ = decoder_lstm(decoder_embedding, initial_state=[state_h, state_c])

# Механизм внимания
attention = BahdanauAttention(lstm_units)
# Получаем маску из энкодера
encoder_mask = encoder_embedding.compute_mask(encoder_inputs)
context_vector = attention(decoder_outputs, encoder_outputs, mask=encoder_mask)

# Объединяем контекстный вектор и выходы декодера
decoder_combined_context = Concatenate(axis=-1)([decoder_outputs, context_vector])

# Полносвязный слой для предсказания
decoder_dense = Dense(answer_vocab_size, activation='softmax', name='decoder_dense')
decoder_outputs = decoder_dense(decoder_combined_context)

# Создание модели
model = Model({'encoder_inputs': encoder_inputs, 'decoder_inputs': decoder_inputs}, decoder_outputs)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.summary()

# 6. Обучение модели
# Преобразование целевых данных в формат, подходящий для loss='sparse_categorical_crossentropy'
decoder_target_data = decoder_target_data.reshape(*decoder_target_data.shape, 1)

# Параметры обучения
epochs = 100
batch_size = 64

# Обучение модели
model.fit(
    x={'encoder_inputs': encoder_input_data, 'decoder_inputs': decoder_input_data},
    y=decoder_target_data,
    batch_size=batch_size,
    epochs=epochs,
    verbose=2
)

# 7. Создание моделей энкодера и декодера для инференса

# Модель энкодера
encoder_model_inf = Model(encoder_inputs, [encoder_outputs, state_h, state_c])

# Порты состояний и скрытых состояний для инференса
decoder_state_input_h = Input(shape=(lstm_units,), name='decoder_state_input_h')
decoder_state_input_c = Input(shape=(lstm_units,), name='decoder_state_input_c')
decoder_hidden_state_input = Input(shape=(max_question_length, lstm_units), name='decoder_hidden_state_input')

# Встраивание декодера
decoder_inf_embedding = decoder_embedding(decoder_inputs)

# Вызов декодера с новыми состояниями
decoder_outputs_inf, state_h_inf, state_c_inf = decoder_lstm(decoder_inf_embedding, initial_state=[decoder_state_input_h, decoder_state_input_c])

# Контекстный вектор с использованием внимания
context_vector_inf = attention(decoder_outputs_inf, decoder_hidden_state_input)
decoder_combined_context_inf = Concatenate(axis=-1)([decoder_outputs_inf, context_vector_inf])

# Предсказание следующего слова
decoder_outputs_inf = decoder_dense(decoder_combined_context_inf)

# Создание модели декодера для инференса
decoder_model_inf = Model(
    [decoder_inputs, decoder_hidden_state_input, decoder_state_input_h, decoder_state_input_c],
    [decoder_outputs_inf, state_h_inf, state_c_inf]
)

# 8. Функция для генерации ответов
def decode_sequence(input_seq):
    # Получение выходов энкодера и состояний
    encoder_outs, state_h, state_c = encoder_model_inf.predict(input_seq)

    # Инициализация целевой последовательности с токеном <START>
    target_seq = np.array([[answer_tokenizer.word_index['<start>']]])

    # Инициализация переменных
    stop_condition = False
    decoded_sentence = ''
    max_dec_steps = max_answer_length

    while not stop_condition:
        output_tokens, h, c = decoder_model_inf.predict(
            [target_seq, encoder_outs, state_h, state_c]
        )

        # Получаем индекс предсказанного слова
        sampled_token_index = np.argmax(output_tokens[0, -1, :])
        sampled_word = answer_tokenizer.index_word.get(sampled_token_index, '')

        # Добавляем предсказанное слово
        if sampled_word != '<end>':
            decoded_sentence += ' ' + sampled_word

        # Проверяем условие остановки
        if (sampled_word == '<end>' or len(decoded_sentence.split()) > max_dec_steps):
            stop_condition = True

        # Обновляем целевую последовательность
        target_seq = np.array([[sampled_token_index]])

        # Обновляем состояния
        state_h, state_c = h, c

    return decoded_sentence.strip()

# 9. Функция для получения ответа на пользовательский вопрос
def get_response(question):
    # Препроцессинг вопроса
    question = question.lower()
    question = question.translate(str.maketrans('', '', string.punctuation))
    tokens = question.split()
    tokens = [word for word in tokens if word not in set(stopwords.words('russian'))]
    clean_question = ' '.join(tokens)

    # Преобразование в последовательность
    seq = question_tokenizer.texts_to_sequences([clean_question])
    padded_seq = pad_sequences(seq, maxlen=max_question_length, padding='post')

    # Генерация ответа
    answer = decode_sequence(padded_seq)
    return answer

# 10. Пример использования модели
user_question = "Чем ты занимаешься?"
response = get_response(user_question)
print("Вопрос:", user_question)
print("Ответ:", response)

# Дополнительный пример
user_question2 = "Откуда ты?"
response2 = get_response(user_question2)
print("\nВопрос:", user_question2)
print("Ответ:", response2)

D:\pythonProject18\.venv\pythonProject\Scripts\python.exe D:\pythonProject\main.py 
2024-11-18 15:26:49.889612: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-11-18 15:26:50.988302: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
[nltk_data] Downloading package stopwords to
[nltk_data]     C:\Users\admin\AppData\Roaming\nltk_data...
[nltk_data]   Package stopwords is already up-to-date!
2024-11-18 15:26:53.561657: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
Препроцессинг завершен.
Вопросы после препроцессинга: ['зовут', 'сколько тебе лет', 'откуда', 'занимаешься', 'сегодня день', 'погода', 'умеешь', 'твой создатель', 'добраться ближайшего магазина', 'можешь рассказать анекдот', 'какое сегодня число', 'какие хобби', 'знаешь машинном обучении', 'можешь помочь домашним заданием', 'какие языки программирования знаешь', 'какова столица франции', 'расскажи космосе', 'такое искусственный интеллект', 'можешь порекомендовать книгу', 'приготовить борщ', 'знаешь python', 'какую музыку предпочитаешь', 'любимый фильм', 'можешь рассказать своей работе', 'каковы твои планы будущее']
Ответы после препроцессинга: ['зовут алекс', '25 лет', 'москвы', 'изучаю машинное обучение', 'сегодня среда', 'сегодня солнечно небольшой облачностью', 'могу вести беседу отвечать вопросы', 'создала команда разработчиков', 'добраться магазина нужно идти прямо улице повернуть направо втором светофоре', 'почему программисты любят природу багов', 'сегодня 18 ноября 2024 года', 'люблю читать программировать', 'машинное обучение — это область искусственного интеллекта которая занимается разработкой алгоритмов позволяющих компьютерам учиться данных', 'нужна помощь', 'знаком python java c другими языками', 'столицей франции является париж', 'космос — это пространство пределами земной атмосферы находятся звезды планеты другие небесные тела', 'искусственный интеллект — это способность машин выполнять задачи требующие человеческого интеллекта', 'рекомендую книгу машинное обучение начинающих', 'приготовления борща понадобятся свекла капуста картофель морковь лук другие ингредиенты сначала обжарьте овощи затем добавьте кипящий бульон варите готовности', 'python — это высокоуровневый язык программирования известный своей простотой широким применением', 'нравится классическая музыка современные электронные жанры', 'любимого фильма могу порекомендовать известные картины', 'создан помогать людям отвечать вопросы вести беседы', 'стремлюсь улучшаться обучаться новым навыкам полезнее пользователей']
Размер словаря вопросов: 56
Размер словаря ответов: 142
Данные подготовлены.
Traceback (most recent call last):
  File "D:\pythonProject\main.py", line 180, in <module>
    encoder_mask = encoder_embedding.compute_mask(encoder_inputs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'KerasTensor' object has no attribute 'compute_mask'