Update app.py

app.py

@@ -1,78 +1,107 @@
 import streamlit as st
-from PyPDF2 import PdfReader
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import pytesseract
+from pdf2image import convert_from_path
+import os
+import re
+import json
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from tqdm import tqdm
 
-# تنظیمات صفحه
-st.set_page_config(page_title="PDF Zusammenfasser", layout="centered")
+# تنظیم Tesseract
+pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'
 
-# بارگذاری مدل و توکنایزر
-@st.cache_resource
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
-    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
-    return tokenizer, model
+# بارگذاری مدل سفارشی
+model_name = "BabakBagheriGisour/NetworkPlus"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-tokenizer, model = load_model()
+# بررسی GPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-# بارگذاری فایل PDF
-uploaded_file = st.file_uploader("Laden Sie Ihre PDF-Datei hoch", type=["pdf"])
+# پاکسازی متن
+def clean_text(text):
+    text = re.sub(r'\s+', ' ', text)  # حذف فضاهای اضافی
+    text = re.sub(r'http\S+|www\.[\w.-]+', '', text)  # حذف لینک‌ها
+    text = re.sub(r'[^\w\sÄäÖöÜüß]+', '', text)  # حذف علائم غیر ضروری
+    return text.strip()
 
-if uploaded_file:
-    # خواندن متن از PDF
-    reader = PdfReader(uploaded_file)
-    pdf_text = ""
-    for page in reader.pages:
-        pdf_text += page.extract_text()
-
-    # نمایش متن کامل
-    st.text_area("Vollständiger Text der PDF-Datei:", pdf_text, height=300)
+# استخراج متن از فایل PDF
+def extract_text_using_ocr(pdf_path):
+    pages = convert_from_path(pdf_path, 300)  # تبدیل PDF به تصاویر
+    all_text = []
+    for page in pages:
+        text = pytesseract.image_to_string(page, lang="deu")  # زبان آلمانی
+        all_text.append(clean_text(text))
+    return all_text
 
-    # محاسبه تعداد خطوط و کلمات
-    lines = pdf_text.split("\n")
-    num_lines = len(lines)
-    num_words = len(pdf_text.split())
+# خلاصه‌سازی با مدل سفارشی
+def summarize_text(text):
+    inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True, padding=True).to(device)
+    summary_ids = model.generate(inputs['input_ids'], num_beams=4, max_length=100, early_stopping=True)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
 
-    # نمایش اطلاعات فایل
-    st.info(f"**Anzahl der Zeilen:** {num_lines}")
-    st.info(f"**Anzahl der Wörter:** {num_words}")
+# پردازش PDF و تولید JSONL
+def process_pdf_to_jsonl(pdf_file, progress_bar):
+    pages_text = extract_text_using_ocr(pdf_file)
+    total_lines = sum(len(page.splitlines()) for page in pages_text)
+    total_words = sum(len(page.split()) for page in pages_text)
+    
+    data = []
+    for idx, page_text in enumerate(pages_text):
+        summary = summarize_text(page_text)
+        data.append({
+            "text": page_text,
+            "zusammenfassen": summary
+        })
+        # به‌روزرسانی نوار پیشرفت
+        progress_bar.progress((idx + 1) / len(pages_text))
+    
+    return data, len(pages_text), total_lines, total_words
 
-    # شناسایی موضوع فایل
-    words = [word.lower() for word in pdf_text.split() if word.isalpha()]
-    topic = max(set(words), key=words.count)  # بیشترین تکرار یک کلمه
-    st.info(f"**Mögliches Thema der Datei:** {topic}")
+# رابط کاربری با Streamlit
+st.title("PDF to JSONL Converter with BabakBagheriGisour/NetworkPlus")
+uploaded_file = st.file_uploader("Bitte laden Sie eine PDF-Datei hoch", type="pdf")
 
-    # پردازش خط به خط برای خلاصه‌سازی
-    st.write("Das Modell verarbeitet den Text...")
+if uploaded_file:
+    # ذخیره فایل آپلود شده
+    temp_file_path = f"temp_{uploaded_file.name}"
+    with open(temp_file_path, "wb") as f:
+        f.write(uploaded_file.read())
+    
+    # پردازش فایل
+    st.info("Das Modell verarbeitet den Text...")
     progress_bar = st.progress(0)
-    final_summary = ""
+    data, total_pages, total_lines, total_words = process_pdf_to_jsonl(temp_file_path, progress_bar)
 
-    for idx, line in enumerate(lines):
-        if len(line.strip()) == 0:  # عبور از خطوط خالی
-            continue
-    
-        # تبدیل خط به توکن
-        input_ids = tokenizer.encode(line, return_tensors="pt", truncation=True, max_length=1024)
-    
-        # تولید خلاصه برای هر خط
-        try:
-            summary_ids = model.generate(
-                input_ids=input_ids,
-                num_beams=2,
-                max_length=50,  # حداکثر طول خ��اصه
-                min_length=10  # حداقل طول خلاصه، در صورت نیاز قابل تنظیم است
-            )
-            decoded_summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-            final_summary += decoded_summary + " "
-        except Exception as e:
-            st.warning(f"Fehler bei Zeile {idx + 1}: {e}")
-    
-        # نمایش خلاصه فعلی
-        st.write(f"**Zeile {idx + 1}:** {decoded_summary}")
-    
-        # به‌روزرسانی نوار پیشرفت
-        progress_bar.progress((idx + 1) / num_lines)
+    # نمایش اطلاعات فایل
+    st.subheader("Dateiinformationen:")
+    st.write(f"**Anzahl der Seiten:** {total_pages}")
+    st.write(f"**Anzahl der Zeilen:** {total_lines}")
+    st.write(f"**Anzahl der Wörter:** {total_words}")
 
+    # نمایش خلاصه‌ها
+    st.subheader("Zusammenfassungen:")
+    for idx, item in enumerate(data):
+        st.write(f"**Seite {idx + 1}:**")
+        st.text_area("Originaltext", item["text"], height=200)
+        st.text_area("Zusammenfassung", item["zusammenfassen"], height=100)
+
+    # ذخیره به JSONL
+    output_file = f"{uploaded_file.name.split('.')[0]}.jsonl"
+    with open(output_file, 'w', encoding='utf-8') as f:
+        for item in data:
+            f.write(json.dumps(item, ensure_ascii=False) + '\n')
+    
+    st.success("Verarbeitung abgeschlossen!")
+    st.download_button(
+        label="Download JSONL",
+        data=open(output_file, "rb").read(),
+        file_name=output_file,
+        mime="application/jsonl"
+    )
 
-    # نمایش خلاصه نهایی
-    st.success("Zusammenfassung der Datei:")
-    st.write(final_summary)
+    # حذف فایل‌های موقت
+    os.remove(temp_file_path)
+    os.remove(output_file)