import aiohttp
import asyncio
import json
import logging
import os
import time
import torch
from flask import Flask, request, jsonify, render_template
from bs4 import BeautifulSoup
from sentence_transformers import SentenceTransformer, util
from transformers import pipeline, RobertaTokenizerFast, RobertaForQuestionAnswering, Trainer, TrainingArguments
from datasets import load_dataset
import evaluate
from razdel import sentenize

# Настройка логирования
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')

app = Flask(__name__)

# Инициализация моделей
logging.info("Загрузка моделей...")
try:
    # Модель для семантического поиска
    search_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
    # Модель для ответов на вопросы (изначально)
    qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
    logging.info("Модели успешно загружены.")
except Exception as e:
    logging.error(f"Ошибка при загрузке моделей: {e}")
    raise e  # Прерываем запуск приложения, если модели не загрузились

# Пути к файлам
ARTICLES_FILE = 'articles.json'
EMBEDDINGS_FILE = 'embeddings.pt'
QA_DATASET_FILE = 'qa_dataset.json'  # Новый файл для QA датасета


# Асинхронные функции для загрузки статей
async def fetch_random_wiki_title(session):
    try:
        async with session.get('https://ru.wikipedia.org/wiki/Special:Random') as response:
            await asyncio.sleep(0.05)  # Пауза между запросами
            return str(response.url)
    except Exception as e:
        logging.error(f"Ошибка при запросе случайной статьи: {e}")
        return None


async def fetch_wiki_article(session, url):
    try:
        async with session.get(url) as response:
            content = await response.text()
            soup = BeautifulSoup(content, 'html.parser')

            title_elem = soup.find(id="firstHeading")
            if not title_elem:
                logging.error(f"Не удалось найти заголовок статьи по URL: {url}")
                return None
            title = title_elem.text.strip()

            paragraphs = [p.get_text() for p in soup.find_all('p') if len(p.get_text()) > 50]

            if not paragraphs:
                logging.warning(f"Нет подходящих параграфов для статьи: {title}")
                return None

            article = ' '.join(paragraphs)
            return {'url': url, 'title': title, 'article': article}
    except Exception as e:
        logging.error(f"Ошибка при запросе статьи по URL {url}: {e}")
        return None


async def get_wiki_articles(count, concurrency=10):
    semaphore = asyncio.Semaphore(concurrency)
    articles = []
    logging.info(f"Запуск загрузки {count} статей с Википедии...")

    async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(limit=concurrency)) as session:
        tasks = [fetch_random_wiki_title(session) for _ in range(count)]
        urls = await asyncio.gather(*tasks)

        fetch_tasks = [fetch_wiki_article(session, url) for url in urls if url]
        articles_data = await asyncio.gather(*fetch_tasks)

    articles = [article for article in articles_data if article]
    if articles:
        try:
            with open(ARTICLES_FILE, 'w', encoding='utf-8') as f:
                json.dump(articles, f, ensure_ascii=False, indent=2)
            logging.info(f"Загрузка завершена. Статьи сохранены в '{ARTICLES_FILE}'.")
        except Exception as e:
            logging.error(f"Ошибка при сохранении статей в '{ARTICLES_FILE}': {e}")
    else:
        logging.warning("Не было загружено ни одной статьи.")

    return articles


# Загрузка и подготовка статей
def load_articles():
    if not os.path.exists(ARTICLES_FILE):
        logging.warning(f"Файл '{ARTICLES_FILE}' не найден.")
        return []
    try:
        with open(ARTICLES_FILE, 'r', encoding='utf-8') as f:
            articles = json.load(f)
            logging.info(f"Загружено {len(articles)} статей из '{ARTICLES_FILE}'.")
            return articles
    except Exception as e:
        logging.error(f"Ошибка при загрузке '{ARTICLES_FILE}': {e}")
        return []


def split_into_chunks(article_text, max_length=512):
    sentences = [sentence.text for sentence in sentenize(article_text)]
    chunks = []
    current_chunk = ""
    for sentence in sentences:
        if len(current_chunk) + len(sentence) + 1 <= max_length:
            current_chunk += " " + sentence
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = sentence
    if current_chunk:
        chunks.append(current_chunk.strip())
    return chunks


def process_articles(articles):
    processed = []
    for article in articles:
        chunks = split_into_chunks(article['article'])
        for chunk in chunks:
            processed.append({
                'url': article['url'],
                'title': article['title'],
                'chunk': chunk
            })
    logging.info(f"Обработано {len(processed)} чанков из статей.")
    return processed


def generate_article_embeddings(processed_articles):
    try:
        chunks = [article['chunk'] for article in processed_articles]
        embeddings = search_model.encode(chunks, convert_to_tensor=True, show_progress_bar=True)
        logging.info("Эмбеддинги статей успешно сгенерированы.")
        return embeddings
    except Exception as e:
        logging.error(f"Ошибка при генерации эмбеддингов: {e}")
        return None


def save_embeddings(embeddings, processed_articles):
    try:
        torch.save({
            'embeddings': embeddings,
            'data': processed_articles
        }, EMBEDDINGS_FILE)
        logging.info(f"Эмбеддинги сохранены в '{EMBEDDINGS_FILE}'.")
    except Exception as e:
        logging.error(f"Ошибка при сохранении эмбеддингов в '{EMBEDDINGS_FILE}': {e}")


def load_or_generate_embeddings():
    if os.path.exists(EMBEDDINGS_FILE):
        logging.info("Загрузка эмбеддингов из файла...")
        try:
            data = torch.load(EMBEDDINGS_FILE, map_location=torch.device('cpu'))
            embeddings = data['embeddings']
            processed_articles = data['data']
            logging.info("Эмбеддинги успешно загружены.")
            return embeddings, processed_articles
        except Exception as e:
            logging.error(f"Ошибка при загрузке эмбеддингов из '{EMBEDDINGS_FILE}': {e}")
            return None, None
    else:
        logging.info("Эмбеддинги не найдены. Начинается их генерация...")
        articles = load_articles()
        if not articles:
            logging.error("Нет доступных статей для генерации эмбеддингов.")
            return None, None
        processed_articles = process_articles(articles)
        embeddings = generate_article_embeddings(processed_articles)
        if embeddings is not None:
            save_embeddings(embeddings, processed_articles)
            return embeddings, processed_articles
        else:
            return None, None


# Поиск релевантных статей на основе вопроса
def find_relevant_articles(question, embeddings, processed_articles, top_k=5):
    try:
        question_embedding = search_model.encode(question, convert_to_tensor=True)
        cos_scores = util.cos_sim(question_embedding, embeddings)[0]
        top_results = torch.topk(cos_scores, k=top_k)

        relevant_articles = []
        for score, idx in zip(top_results.values, top_results.indices):
            if score.item() > 0.1:  # Порог релевантности
                relevant_articles.append(processed_articles[idx])
        logging.info(f"Найдено {len(relevant_articles)} релевантных статей.")
        return relevant_articles
    except Exception as e:
        logging.error(f"Ошибка при поиске релевантных статей: {e}")
        return []


# Генерация ответа на основе релевантных статей
def generate_answer(question, relevant_articles):
    answers = []
    for article in relevant_articles:
        try:
            result = qa_pipeline({
                'question': question,
                'context': article['chunk']
            })
            answers.append((result['score'], result['answer']))
        except Exception as e:
            logging.error(f"Ошибка при генерации ответа из статьи '{article['title']}': {e}")

    if not answers:
        logging.info("Ответы не найдены.")
        return "Ответ не найден."

    # Выбор ответа с наивысшим доверием
    best_answer = max(answers, key=lambda x: x[0])
    if best_answer[0] > 0.1:
        logging.info("Лучший ответ найден.")
        return best_answer[1]
    else:
        logging.info("Нет достаточного доверия к найденным ответам.")
        return "Ответ не найден."


# Загрузка QA датасета
def load_qa_dataset(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            qa_data = json.load(f)
        logging.info(f"QA датасет загружен из '{file_path}'.")
        return qa_data
    except Exception as e:
        logging.error(f"Ошибка при загрузке QA датасета из '{file_path}': {e}")
        return None


# Функция для fine-tuning модели QA
def fine_tune_qa_model(qa_data):
    # Загрузка датасета из JSON
    dataset = load_dataset('json', data_files={'train': QA_DATASET_FILE}, field='data')

    # Разделение на тренировочный и валидационный наборы
    split_dataset = dataset['train'].train_test_split(test_size=0.1, seed=42)
    train_dataset = split_dataset['train']
    eval_dataset = split_dataset['test']

    # Загрузка токенизатора и модели
    tokenizer = RobertaTokenizerFast.from_pretrained('deepset/roberta-base-squad2')
    model = RobertaForQuestionAnswering.from_pretrained('deepset/roberta-base-squad2')

    # Исправленная функция предварительной обработки
    def preprocess_function(examples):
        questions = []
        contexts = []
        answers = []

        for article_paragraphs in examples['paragraphs']:
            for paragraph in article_paragraphs:
                context = paragraph['context']
                for qa in paragraph['qas']:
                    questions.append(qa['question'])
                    contexts.append(context)
                    answers.append(qa['answers'][0])

        encodings = tokenizer(
            questions,
            contexts,
            truncation=True,
            padding=True,
            max_length=512,
            return_offsets_mapping=True
        )

        start_positions = []
        end_positions = []

        for i, answer in enumerate(answers):
            answer_text = answer['text']
            answer_start_char = answer['answer_start']
            answer_end_char = answer_start_char + len(answer_text)

            offsets = encodings['offset_mapping'][i]
            start_token = None
            end_token = None

            for idx, (start, end) in enumerate(offsets):
                if start <= answer_start_char < end:
                    start_token = idx
                if start < answer_end_char <= end:
                    end_token = idx
                    break

            if start_token is None:
                start_token = 0
            if end_token is None:
                end_token = len(offsets) - 1

            start_positions.append(start_token)
            end_positions.append(end_token)

        encodings.update({'start_positions': start_positions, 'end_positions': end_positions})
        return encodings

    # Токенизация данных
    tokenized_train = train_dataset.map(
        preprocess_function,
        batched=True,
        remove_columns=['title', 'paragraphs']
    )
    tokenized_eval = eval_dataset.map(
        preprocess_function,
        batched=True,
        remove_columns=['title', 'paragraphs']
    )

    # Определение аргументов для обучения
    training_args = TrainingArguments(
        output_dir='./results',
        evaluation_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,
        num_train_epochs=3,
        weight_decay=0.01,
    )

    # Загрузка метрики SQuAD
    metric = evaluate.load("squad")

    def compute_metrics(eval_pred):
        predictions, references = eval_pred
        return metric.compute(predictions=predictions, references=references)

    # Создание Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_train,
        eval_dataset=tokenized_eval,
        tokenizer=tokenizer,
        compute_metrics=compute_metrics
    )

    # Запуск обучения
    trainer.train()

    # Сохранение модели
    model.save_pretrained('fine_tuned_roberta_qa')
    tokenizer.save_pretrained('fine_tuned_roberta_qa')
    logging.info("Модель QA успешно обучена и сохранена.")


# Маршрут главной страницы
@app.route('/')
def index():
    return render_template('index.html')


# Маршрут для загрузки статей
@app.route('/getwiki', methods=['POST'])
def get_wiki():
    data = request.get_json()
    count = data.get('count', 1)

    if not isinstance(count, int) or count <= 0:
        return jsonify({"error": "Параметр 'count' должен быть положительным целым числом."}), 400

    # Замер времени начала
    start_time = time.time()

    # Запуск асинхронной загрузки статей
    try:
        asyncio.run(get_wiki_articles(count))
    except Exception as e:
        logging.error(f"Ошибка при загрузке статей: {e}")
        return jsonify({"error": "Ошибка при загрузке статей."}), 500

    # Замер времени окончания
    duration = time.time() - start_time

    # После загрузки обновить эмбеддинги
    embeddings, processed_articles = load_or_generate_embeddings()
    if embeddings is None:
        return jsonify({"error": "Не удалось загрузить или сгенерировать эмбеддинги."}), 500

    return jsonify({
        "message": f"Загружено {count} статей.",
        "execution_time": f"{duration:.2f} секунд"
    })


# Маршрут для обработки вопросов
@app.route('/ask', methods=['POST'])
def ask():
    data = request.get_json()
    question = data.get('question', '').strip()

    if not question:
        return jsonify({"error": "Введите вопрос."}), 400

    # Загрузка эмбеддингов и обработанных статей
    embeddings, processed_articles = load_or_generate_embeddings()
    if embeddings is None or processed_articles is None:
        return jsonify({"error": "Нет эмбеддингов. Сначала загрузите статьи."}), 400

    # Поиск релевантных статей
    relevant_articles = find_relevant_articles(question, embeddings, processed_articles, top_k=5)

    if not relevant_articles:
        return jsonify({"answer": "Не найдено подходящих статей.", "execution_time": "0.00 секунд"})

    # Замер времени генерации ответа
    start_time = time.time()

    # Генерация ответа
    answer = generate_answer(question, relevant_articles)

    # Замер времени окончания
    duration = time.time() - start_time

    return jsonify({
        "answer": answer,
        "execution_time": f"{duration:.2f} секунд"
    })


# Маршрут для обучения модели QA
@app.route('/train_qa', methods=['POST'])
def train_qa():
    # Загружаем QA датасет
    qa_data = load_qa_dataset(QA_DATASET_FILE)
    if not qa_data:
        return jsonify({"error": "Не удалось загрузить QA датасет."}), 500

    # Запуск обучения
    try:
        fine_tune_qa_model(qa_data)
    except Exception as e:
        logging.error(f"Ошибка при обучении модели QA: {e}")
        return jsonify({"error": "Ошибка при обучении модели QA."}), 500

    return jsonify({"message": "Модель QA успешно обучена и сохранена."})


if __name__ == '__main__':
    # Проверка наличия начальных данных
    if not os.path.exists(ARTICLES_FILE):
        logging.warning(f"'{ARTICLES_FILE}' не найден. Рекомендуется загрузить статьи через '/getwiki'.")
    if not os.path.exists(EMBEDDINGS_FILE):
        logging.warning(f"'{EMBEDDINGS_FILE}' не найден. Эмбеддинги будут сгенерированы при первой загрузке статей.")
    if not os.path.exists(QA_DATASET_FILE):
        logging.warning(
            f"'{QA_DATASET_FILE}' не найден. Рекомендуется создать QA датасет для обучения модели через '/train_qa'.")

    app.run(host='0.0.0.0', port=5000)
 * Running on http://127.0.0.1:5000
 * Running on http://192.168.1.3:5000

ошибки при обучении
2024-10-15 20:57:23,956 INFO Press CTRL+C to quit
2024-10-15 20:57:26,698 INFO QA датасет загружен из 'qa_dataset.json'.
Generating train split: 1 examples [00:00, 142.74 examples/s]
2024-10-15 20:57:27,806 ERROR Ошибка при обучении модели QA: With n_samples=1, test_size=0.1 and train_size=None, the resulting train set will be empty. Adjust any of the aforementioned parameters.
2024-10-15 20:57:27,806 INFO 127.0.0.1 - - [15/Oct/2024 20:57:27] "POST /train_qa HTTP/1.1" 500 -
2024-10-15 21:02:48,797 INFO QA датасет загружен из 'qa_dataset.json'.
2024-10-15 21:02:49,627 ERROR Ошибка при обучении модели QA: With n_samples=1, test_size=0.1 and train_size=None, the resulting train set will be empty. Adjust any of the aforementioned parameters.
2024-10-15 21:02:49,627 INFO 127.0.0.1 - - [15/Oct/2024 21:02:49] "POST /train_qa HTTP/1.1" 500 -