Update app.py

app.py

@@ -3,7 +3,21 @@ from transformers import pipeline
 import pdfplumber
 import os
 
-# --- Ленивая загрузка модели ---
+# OCR
+try:
+    import pytesseract
+    from pdf2image import convert_from_path
+    from PIL import Image
+    OCR_AVAILABLE = True
+except ImportError:
+    OCR_AVAILABLE = False
+
+# DOCX / PDF экспорт
+from docx import Document
+from fpdf import FPDF
+
+# ---------- МОДЕЛЬ ----------
+
 summarizer = None
 
 def load_model():
@@ -15,9 +29,13 @@ def load_model():
         )
     return summarizer
 
-# --- Безопасное чтение PDF ---
-def extract_pdf_text(path):
+# ---------- ЧТЕНИЕ ФАЙЛА ----------
+
+def extract_pdf_text(path: str):
+    """Пытаемся вытащить текст из PDF. Если текста нет – пробуем OCR."""
     text = ""
+
+    # 1) обычный текстовый PDF
     try:
         with pdfplumber.open(path) as pdf:
             for page in pdf.pages:
@@ -26,59 +44,144 @@ def extract_pdf_text(path):
                     text += chunk + "\n"
     except Exception as e:
         return "", f"Ошибка при чтении PDF: {e}"
-    return text, None
 
-# --- Универсальное чтение файла ---
-def read_file(path):
+    if text.strip():
+        return text, None
+
+    # 2) если текст не найден – пробуем OCR
+    if not OCR_AVAILABLE:
+        return "", "PDF выглядит как скан (изображение). Для OCR нужно pytesseract + tesseract-ocr."
+
+    try:
+        images = convert_from_path(path, dpi=200)
+        ocr_text = ""
+        for img in images:
+            ocr_text += pytesseract.image_to_string(img) + "\n"
+        if not ocr_text.strip():
+            return "", "OCR не смог распознать текст в этом PDF."
+        return ocr_text, None
+    except Exception as e:
+        return "", f"Ошибка OCR при обработке PDF: {e}"
+
+
+def read_file(path: str):
     if not path:
         return "", "Файл не передан."
+
     path = str(path).strip()
+    lower = path.lower()
 
-    if path.lower().endswith(".pdf"):
+    if lower.endswith(".pdf"):
         return extract_pdf_text(path)
 
+    # TXT / другие текстовые файлы
     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):
+# ---------- ЧАНКИНГ ТЕКСТА ----------
+
+def chunk_text(text: str, max_chars: int = 2500):
+    """Режем длинный текст на куски, стараясь обрезать по точкам."""
+    chunks = []
+    while len(text) > max_chars:
+        cut = text[:max_chars]
+        last_dot = cut.rfind(".")
+        if last_dot != -1:
+            cut = cut[:last_dot + 1]
+        chunks.append(cut)
+        text = text[len(cut):]
+    chunks.append(text)
+    return chunks
+
+# ---------- СОХРАНЕНИЕ РЕЗЮМЕ В DOCX/PDF ----------
+
+def save_docx(summary: str) -> str:
+    doc = Document()
+    doc.add_heading("Eroha Summarizer – Резюме документа", level=1)
+    for paragraph in summary.split("\n"):
+        doc.add_paragraph(paragraph)
+    path = "/tmp/summary.docx"
+    doc.save(path)
+    return path
+
+def save_pdf(summary: str) -> str:
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.set_font("Arial", size=12)
+
+    for line in summary.split("\n"):
+        # multi_cell сам переносит строки
+        pdf.multi_cell(0, 8, line)
+    path = "/tmp/summary.pdf"
+    pdf.output(path)
+    return path
+
+# ---------- ОСНОВНАЯ ФУНКЦИЯ ----------
+
+def summarize_file(path: str):
     text, err = read_file(path)
     if err:
-        return f"⚠️ {err}"
+        return f"⚠️ {err}", None, None
 
     if not text.strip():
-        return "⚠️ Не удалось извлечь текст."
-
-    if len(text) < 80:
-        return "⚠️ Слишком мало текста для суммаризации."
+        return "⚠️ Не удалось извлечь текст из файла.", None, None
 
+    # Модель
     model = load_model()
 
-    # Ограничиваем текст для избежания OOM
-    text = text[:4000]
+    # Чанкинг
+    chunks = chunk_text(text, max_chars=2500)
+    partial_summaries = []
 
-    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), и модель создаст краткое резюме."
-)
+    for chunk in chunks:
+        if not chunk.strip():
+            continue
+        try:
+            summary = model(
+                chunk,
+                max_length=180,
+                min_length=60,
+                do_sample=False
+            )
+            partial_summaries.append(summary[0]["summary_text"])
+        except Exception as e:
+            partial_summaries.append(f"[Ошибка в блоке суммаризации: {e}]")
+
+    if not partial_summaries:
+        return "⚠️ Не удалось создать резюме.", None, None
+
+    final_summary = "\n\n".join(partial_summaries)
+
+    # Файлы экспорта
+    docx_path = save_docx(final_summary)
+    pdf_path = save_pdf(final_summary)
+
+    return final_summary, docx_path, pdf_path
+
+# ---------- ИНТЕРФЕЙС GRADIO ----------
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Eroha Summarizer 🧠")
+    gr.Markdown("Загрузите документ (PDF или TXT), и модель создаст краткое резюме с возможностью скачивания DOCX и PDF.")
+
+    with gr.Row():
+        file_input = gr.File(type="filepath", label="Загрузите файл (.pdf или .txt)")
+        with gr.Column():
+            summary_output = gr.Textbox(label="Результат суммаризации", lines=20)
+            docx_output = gr.File(label="Скачать DOCX")
+            pdf_output = gr.File(label="Скачать PDF")
+
+    submit_btn = gr.Button("Запустить суммаризацию")
+
+    submit_btn.click(
+        fn=summarize_file,
+        inputs=file_input,
+        outputs=[summary_output, docx_output, pdf_output]
+    )
 
 if __name__ == "__main__":
     demo.launch()