eroha-agentapi / app.py
Yermek68's picture
Update app.py
0dfa820 verified
raw
history blame
2.46 kB
import gradio as gr
from transformers import pipeline
import pdfplumber
import os
# --- Ленивая загрузка модели ---
summarizer = None
def load_model():
global summarizer
if summarizer is None:
summarizer = pipeline(
"summarization",
model="sshleifer/distilbart-cnn-12-6"
)
return summarizer
# --- Безопасное чтение PDF ---
def extract_pdf_text(path):
text = ""
try:
with pdfplumber.open(path) as pdf:
for page in pdf.pages:
chunk = page.extract_text()
if chunk:
text += chunk + "\n"
except Exception as e:
return "", f"Ошибка при чтении PDF: {e}"
return text, None
# --- Универсальное чтение файла ---
def read_file(path):
if not path:
return "", "Файл не передан."
path = str(path).strip()
if path.lower().endswith(".pdf"):
return extract_pdf_text(path)
try:
with open(path, "r", encoding="utf-8", errors="ignore") as f:
return f.read(), None
except Exception as e:
return "", f"Ошибка при чтении TXT: {e}"
# --- Основная функция ---
def summarize_file(path):
text, err = read_file(path)
if err:
return f"⚠️ {err}"
if not text.strip():
return "⚠️ Не удалось извлечь текст."
if len(text) < 80:
return "⚠️ Слишком мало текста для суммаризации."
model = load_model()
# Ограничиваем текст для избежания OOM
text = text[:4000]
try:
summary = model(
text,
max_length=180,
min_length=60,
do_sample=False
)
return summary[0]["summary_text"]
except Exception as e:
return f"⚠️ Ошибка суммаризации: {e}"
# --- Интерфейс Gradio ---
demo = gr.Interface(
fn=summarize_file,
inputs=gr.File(type="filepath", label="Загрузите файл (.pdf или .txt)"),
outputs=gr.Textbox(label="Результат суммаризации"),
title="Eroha Summarizer 🧠",
description="Загрузите документ (PDF или TXT), и модель создаст краткое резюме."
)
if __name__ == "__main__":
demo.launch()