app.py

import gradio as gr
import glob
from docx import Document
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

def is_header(txt):
    # Абсолютно короткая фраза без знака препинания и вся в верхнем регистре — заголовок
    if not txt or len(txt) < 35:
        if txt == txt.upper() and not txt.endswith(('.', ':', '?', '!')):
            return True
        # Также часто заголовок — просто пара слов с заглавных (мало слов и нет в конце точки):
        if txt.istitle() and len(txt.split()) < 6 and not txt.endswith(('.', ':', '?', '!')):
            return True
    return False

def get_blocks_from_docx():
    docx_list = glob.glob("*.docx")
    if not docx_list:
        return [], []
    doc = Document(docx_list[0])
    blocks = []
    non_header_blocks = []
    for p in doc.paragraphs:
        txt = p.text.strip()
        if (
            txt
            and not (len(txt) <= 3 and txt.isdigit())
            and len(txt.split()) > 3
        ):
            blocks.append(txt)
            if not is_header(txt) and len(txt) > 25:
                non_header_blocks.append(txt)
    # Таблицы
    for table in doc.tables:
        for row in table.rows:
            row_text = " | ".join(cell.text.strip() for cell in row.cells if cell.text.strip())
            if row_text and len(row_text.split()) > 3 and len(row_text) > 25:
                blocks.append(row_text)
                if not is_header(row_text):
                    non_header_blocks.append(row_text)
    # Убираем дубли
    seen = set()
    blocks_clean = []
    non_hdr_clean = []
    for b in blocks:
        if b not in seen:
            blocks_clean.append(b)
            seen.add(b)
    seen = set()
    for b in non_header_blocks:
        if b not in seen:
            non_hdr_clean.append(b)
            seen.add(b)
    return blocks_clean, non_hdr_clean

blocks, normal_blocks = get_blocks_from_docx()
if not blocks or not normal_blocks:
    # Если ничего не нашли — фэйк заглушка
    blocks = ["База знаний пуста: проверьте содержание и формат вашего .docx!"]
    normal_blocks = ["База знаний пуста: проверьте содержание и формат вашего .docx!"]

vectorizer = TfidfVectorizer(lowercase=True).fit(blocks)
matrix = vectorizer.transform(blocks)

tokenizer = T5Tokenizer.from_pretrained("cointegrated/rut5-base-multitask")
model = T5ForConditionalGeneration.from_pretrained("cointegrated/rut5-base-multitask")
model.eval()
device = 'cpu'

def rut5_answer(question, context):
    prompt = f"question: {question} context: {context}"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_length=250, num_beams=4, min_length=40,
            no_repeat_ngram_size=3, do_sample=False
        )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

def ask_chatbot(question):
    question = question.strip()
    if not question:
        return "Пожалуйста, введите вопрос."
    if not normal_blocks or normal_blocks == ["База знаний пуста: проверьте содержание и формат вашего .docx!"]:
        return "Ошибка: база знаний пуста. Проверьте .docx и перезапустите Space."

    user_vec = vectorizer.transform([question.lower()])
    sims = cosine_similarity(user_vec, matrix)
    n_blocks = min(3, len(blocks))
    if n_blocks == 0:
        return "База знаний пуста: загрузите методичку с осмысленными абзацами!"
    # Получаем индексы лучших блоков среди ВСЕХ
    top_idxs = list(reversed(sims.argsort()[-n_blocks:]))
    # Для генерации контекста используем все блоки, но...
    context_blocks = []
    for rank, idx in enumerate(top_idxs):
        if 0 <= idx < len(blocks):
            context_blocks.append(blocks[idx])
    context = " ".join(context_blocks)
    # ...для финального ответа ищем САМЫЙ релевантный не-заголовок (абзац)! 
    # (обычно первый релевантен)
    best_normal_block = ""
    max_sim = -1
    for idx, nb in enumerate(normal_blocks):
        v_nb = vectorizer.transform([nb.lower()])
        sim = cosine_similarity(user_vec, v_nb)[0]
        if sim > max_sim:
            max_sim = sim
            best_normal_block = nb
    # Если совсем всё плохо — fallback на обычный context
    if not best_normal_block:
        best_normal_block = context_blocks if context_blocks else ""
    # Генерируем развернутый ответ с подложкой из максимального контекста
    answer = rut5_answer(question, context)
    # Если слишком кратко — дублируем релевантный фрагмент (абзац)
    if len(answer.strip().split()) < 8 or answer.count('.') < 2:
        answer += "\n\n" + best_normal_block
    # Финальный ответ — если сгенерированный ответ случайно "превратился" в заголовок, заменяем его на абзац!
    if is_header(answer):
        answer = best_normal_block
    return answer

EXAMPLES = [
    "Какие требования к объему магистерской диссертации?",
    "Как оформить список литературы?",
    "Какие сроки сдачи и защиты ВКР?",
    "Что должно быть во введении?",
    "Какой процент оригинальности требуется?",
    "Как оформлять формулы?"
]

with gr.Blocks() as demo:
    gr.Markdown(
        "# Русскоязычный FAQ-чат-бот по методичке (AI+документ)\nЗадайте вопрос — получите развернутый ответ на основании вашего документа!"
    )
    question = gr.Textbox(label="Ваш вопрос", lines=2)
    ask_btn = gr.Button("Получить ответ")
    answer = gr.Markdown(label="Ответ", visible=True)

    def with_spinner(q):
        yield "Чат-бот думает..."
        yield ask_chatbot(q)

    ask_btn.click(with_spinner, question, answer)
    question.submit(with_spinner, question, answer)
    gr.Markdown("#### Примеры вопросов:")
    gr.Examples(EXAMPLES, inputs=question)
    gr.Markdown("""
    ---
    ### Контакты (укажите свои)
    Преподаватель: ___________________
    Email: ___________________________
    Кафедра: _________________________
    """)

demo.launch()