улучши
import os
import json
import numpy as np
from collections import Counter
from sklearn.model_selection import train_test_split
from transformers import AutoTokenizer, TFAutoModel
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from tensorflow.keras.layers import Input, Dense, Dropout
from tensorflow.keras.models import Model
import matplotlib.pyplot as plt

# Отключение oneDNN (по желанию)
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

# Проверка наличия GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(f"{len(gpus)} Physical GPUs, {len(logical_gpus)} Logical GPUs")
    except RuntimeError as e:
        print(e)
else:
    print("No GPU available, using CPU.")


def load_and_filter_data(json_file_path, top_n_answers=1000, num_samples=None):
    """
    Загружает данные из JSON файла, фильтрует по частоте ответов и выбирает заданное количество примеров.

    Parameters:
    - json_file_path: Путь к файлу с данными.
    - top_n_answers: Количество наиболее часто встречающихся ответов для сохранения.
    - num_samples: Количество примеров для использования (если None, используются все).

    Returns:
    - filtered_training_data: Список словарей с отфильтрованными данными.
    - output_map: Маппинг ответов на индексы.
    """
    with open(json_file_path, 'r', encoding='utf-8') as file:
        training_data = json.load(file)

    # Получение уникальных ответов и их частот
    answer_counter = Counter(item['answer'] for item in training_data)
    common_answers = answer_counter.most_common(top_n_answers)  # Топ-N самых частых ответов

    # Фильтрация данных, оставляя только общие ответы
    filtered_training_data = [item for item in training_data if item['answer'] in dict(common_answers)]

    if num_samples is not None and num_samples < len(filtered_training_data):
        filtered_training_data = filtered_training_data[:num_samples]

    # Получение уникальных ответов и создание маппинга
    output_set = set(item['answer'] for item in filtered_training_data)
    output_array = list(output_set)
    output_map = {output: idx for idx, output in enumerate(output_array)}

    return filtered_training_data, output_map


def texts_to_vectors(tokenizer, texts, max_length=512, batch_size=32):
    """
    Преобразует список текстов в векторные представления с помощью предобученной модели BERT.

    Parameters:
    - tokenizer: Токенизатор BERT.
    - texts: Список текстов для токенизации.
    - max_length: Максимальная длина последовательности.
    - batch_size: Размер батча для преобразования.

    Returns:
    - Векторные представления текстов.
    """
    all_vectors = []
    num_texts = len(texts)
    for i in range(0, num_texts, batch_size):
        batch_texts = texts[i:i + batch_size]
        inputs = tokenizer(batch_texts, return_tensors="tf", padding=True, truncation=True, max_length=max_length)
        outputs = bert_model(**inputs)
        vectors = outputs.last_hidden_state[:, 0, :].numpy()
        all_vectors.append(vectors)
    return np.concatenate(all_vectors, axis=0)


def create_model(input_size, output_size):
    """
    Создает модель нейронной сети поверх BERT.

    Parameters:
    - input_size: Размер входного вектора.
    - output_size: Количество классов (ответов).

    Returns:
    - Компилированная модель Keras.
    """
    inputs = Input(shape=(input_size,))
    x = Dense(512, activation='relu')(inputs)
    x = Dropout(0.2)(x)
    x = Dense(256, activation='relu')(x)
    x = Dropout(0.2)(x)
    x = Dense(128, activation='relu')(x)
    x = Dropout(0.2)(x)
    x = Dense(64, activation='relu')(x)
    x = Dropout(0.2)(x)
    outputs = Dense(output_size, activation='softmax')(x)
    model = Model(inputs=inputs, outputs=outputs)
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model


def prepare_data_for_training(filtered_training_data, output_map, tokenizer):
    """
    Подготавливает данные для обучения: токенизация вопросов и кодирование ответов.

    Parameters:
    - filtered_training_data: Список словарей с данными.
    - output_map: Маппинг ответов на индексы.
    - tokenizer: Токенизатор BERT.

    Returns:
    - input_data: Входные данные для модели (векторные представления вопросов).
    - encoded_output_data: Выходные данные для модели (индексированные ответы).
    """
    questions = [item['question'] for item in filtered_training_data]
    answers = [item['answer'] for item in filtered_training_data]

    input_data = texts_to_vectors(tokenizer, questions)
    encoded_output_data = np.array([output_map[answer] for answer in answers])

    return input_data, encoded_output_data


def train_model(model, X_train, y_train, X_test, y_test, epochs=1000, batch_size=64, patience=5,
                save_path='best_model.h5'):
    """
    Обучает модель и следит за её производительностью с использованием ранней остановки.

    Parameters:
    - model: Компилированная модель Keras.
    - X_train: Входные данные для обучения.
    - y_train: Выходные данные для обучения.
    - X_test: Входные данные для валидации.
    - y_test: Выходные данные для валидации.
    - epochs: Максимальное количество эпох.
    - batch_size: Размер батча.
    - patience: Количество эпох без улучшений для ранней остановки.
    - save_path: Путь для сохранения лучшей модели.

    Returns:
    - history: История обучения модели.
    """
    early_stopping = EarlyStopping(monitor='val_loss', patience=patience, restore_best_weights=True)
    model_checkpoint = ModelCheckpoint(save_path, monitor='val_accuracy', save_best_only=True, mode='max', verbose=1)
    history = model.fit(
        X_train, y_train,
        epochs=epochs,
        batch_size=batch_size,
        validation_data=(X_test, y_test),
        callbacks=[early_stopping, model_checkpoint],
        verbose=1
    )
    return history


def plot_history(history):
    """
    Визуализирует историю обучения модели.

    Parameters:
    - history: История обучения модели.
    """
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']

    epochs_range = range(1, len(acc) + 1)

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(epochs_range, acc, label='Training Accuracy')
    plt.plot(epochs_range, val_acc, label='Validation Accuracy')
    plt.legend(loc='lower right')
    plt.title('Training and Validation Accuracy')

    plt.subplot(1, 2, 2)
    plt.plot(epochs_range, loss, label='Training Loss')
    plt.plot(epochs_range, val_loss, label='Validation Loss')
    plt.legend(loc='upper right')
    plt.title('Training and Validation Loss')

    plt.show()


def predict(model, question, tokenizer, output_array):
    """
    Предсказывает ответ на вопрос с использованием обученной модели.

    Parameters:
    - model: Обученная модель Keras.
    - question: Вопрос для предсказания.
    - tokenizer: Токенизатор BERT.
    - output_array: Список возможных ответов.

    Returns:
    - Предсказанный ответ.
    """
    input_vector = texts_to_vectors(tokenizer, [question]).flatten()
    prediction = model.predict(np.array([input_vector]))
    predicted_class = np.argmax(prediction, axis=-1)[0]
    answer = output_array[predicted_class]
    return answer


def ask_question(model, tokenizer, output_array):
    """
    Интерактивный режим для задавания вопросов и получения ответов от модели.

    Parameters:
    - model: Обученная модель Keras.
    - tokenizer: Токенизатор BERT.
    - output_array: Список возможных ответов.
    """
    print("\nНачало интерактивного режима...")
    print('Введите ваш вопрос (или "exit" для выхода):')
    while True:
        question = input('> ').strip()
        if question.lower() == 'exit':
            print("Выход из программы.")
            break
        try:
            answer = predict(model, question, tokenizer, output_array)
            print(f'Ответ: {answer}\n')
        except Exception as e:
            print(f"Произошла ошибка при обработке вашего запроса: {e}")
            continue


# Главная функция
def main():
    json_file_path = 'bd.json'
    top_n_answers = 1000
    num_samples = int(input(
        "Введите количество примеров для обучения (или нажмите Enter для всех): ")) or None  # Введите число или нажмите Enter для всех

    # Шаг 1: Learn
    print("Шаг 1: Учимся...")
    filtered_training_data, output_map = load_and_filter_data(json_file_path, top_n_answers, num_samples)
    print(f"Загружено {len(filtered_training_data)} примеров.")

    # Загрузка модели BERT
    model_name = "bert-base-cased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    global bert_model
    bert_model = TFAutoModel.from_pretrained(model_name)

    # Подготовка данных
    input_data, encoded_output_data = prepare_data_for_training(filtered_training_data, output_map, tokenizer)

    # Разделение данных на обучающую и тестовую выборки
    X_train, X_test, y_train, y_test = train_test_split(
        input_data,
        encoded_output_data,
        test_size=0.2,
        random_state=42
    )

    print(f"Размер обучающей выборки: {X_train.shape[0]}")
    print(f"Размер тестовой выборки: {X_test.shape[0]}")

    # Шаг 2: Infer
    print("Шаг 2: Инферируем...")
    # Создание модели
    model = create_model(input_data.shape[1], len(output_map))
    print('Модель создана.')

    # Шаг 3: Measure
    print("Шаг 3: Измеряем...")
    # Обучение модели
    history = train_model(model, X_train, y_train, X_test, y_test)
    print('Модель обучена.')

    # Визуализация истории обучения
    plot_history(history)

    # Получение списка возможных ответов
    output_array = list(output_map.keys())

    # Шаг 4: Adjust
    print("Шаг 4: Подстраиваемся...")
    # Если необходимо, можно добавить здесь дополнительные шаги подстройки модели

    # Интерактивный режим
    ask_question(model, tokenizer, output_array)


if __name__ == "__main__":
    main()