TeacherChat / app.py
Wplotnikow's picture
Create app.py
14a84c2 verified
raw
history blame
4.59 kB
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
# 1. Получение текстов из абзацев и таблиц .docx
def get_blocks_from_docx():
docx_list = glob.glob("*.docx")
if not docx_list:
return ["Файл .docx не найден!"]
doc = Document(docx_list[0])
blocks = []
# Абзацы
for p in doc.paragraphs:
txt = p.text.strip()
if txt and not (len(txt) <= 3 and txt.isdigit()):
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:
blocks.append(row_text)
# Удаляем дубли
seen = set()
uniq_blocks = []
for b in blocks:
if b not in seen:
uniq_blocks.append(b)
seen.add(b)
return uniq_blocks
blocks = get_blocks_from_docx()
vectorizer = TfidfVectorizer().fit(blocks)
matrix = vectorizer.transform(blocks)
# 2. Загрузка модели rut5-base-multitask
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=200, num_beams=4)
return tokenizer.decode(output_ids[0], skip_special_tokens=True)
# 3. Комбинированная функция: поиск + генерация
def ask_chatbot(question):
if not question.strip():
return "Пожалуйста, введите вопрос."
if len(blocks) < 2:
return "Ошибка: база знаний пуста или слишком мала. Проверьте .docx."
# Находим релевантный абзац
user_vec = vectorizer.transform([question])
sims = cosine_similarity(user_vec, matrix)[0]
best_idx = sims.argmax()
best_block = blocks[best_idx]
score = sims[best_idx]
if score < 0.12:
context = ""
else:
context = best_block
# Генерируем нейросетевой ответ на русском с учетом найденного контекста
answer = rut5_answer(question, context)
# Для большей прозрачности покажем также фрагмент из документа (можно убрать)
if context:
return f"**Ответ:** {answer}\n\n---\n**Релевантный фрагмент из документа:**\n{context}"
else:
return f"**Ответ:** {answer}\n\n(Контекст в документе не найден — ответ дан на основе общего знания модели.)"
EXAMPLES = [
"Какие требования к объему магистерской диссертации?",
"Как оформить список литературы?",
"Какие сроки сдачи и защиты ВКР?",
"Что должно быть во введении?",
"Какой процент оригинальности требуется?",
"Как оформлять формулы?"
]
with gr.Blocks() as demo:
gr.Markdown(
"""
# Русскоязычный FAQ-чат-бот на базе вашей методички и нейросетевой модели
Задайте вопрос — получайте свежий AI-ответ, опирающийся на ваш документ!
"""
)
question = gr.Textbox(label="Ваш вопрос", lines=2)
ask_btn = gr.Button("Получить ответ")
answer = gr.Markdown(label="Ответ", visible=True)
ask_btn.click(ask_chatbot, question, answer)
question.submit(ask_chatbot, question, answer)
gr.Markdown("#### Примеры вопросов:")
gr.Examples(EXAMPLES, inputs=question)
gr.Markdown("""
---
### Контакты (укажите свои)
Преподаватель: ___________________
Email: ___________________________
Кафедра: _________________________
""")
demo.launch()