Update app.py

app.py

@@ -1,5 +1,7 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import fitz  # PyMuPDF PDF olvasáshoz
+from docx import Document  # DOCX olvasáshoz
 
 # Fordító modell betöltése (M2M100 fordítás magyar és angol között)
 translator_tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
@@ -20,19 +22,41 @@ def translate_to_hungarian(text):
     translated_text = translation_pipeline(text, src_lang="en", tgt_lang="hu", max_length=512)[0]['translation_text']
     return translated_text
 
+# PDF szöveg kinyerése
+def extract_text_from_pdf(pdf_file):
+    text = ""
+    with fitz.open("pdf", pdf_file.read()) as doc:
+        for page in doc:
+            text += page.get_text()
+    return text
+
+# DOCX szöveg kinyerése
+def extract_text_from_docx(docx_file):
+    doc = Document(docx_file)
+    text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
+    return text
+
 # Streamlit alkalmazás
 st.title("AI Jogi Chatbot - Fordítás Alapú Megközelítéssel")
 
 # Dokumentum feltöltése
-uploaded_file = st.file_uploader("Töltsön fel egy dokumentumot", type=["txt"])
+uploaded_file = st.file_uploader("Töltsön fel egy dokumentumot", type=["txt", "pdf", "docx"])
 
 if uploaded_file:
-    try:
-        # Szöveg beolvasása
-        context = uploaded_file.read().decode("utf-8")
-    except UnicodeDecodeError:
-        context = uploaded_file.read().decode("ISO-8859-1")
-    
+    file_type = uploaded_file.type
+
+    # Dokumentum típus szerinti beolvasás
+    if file_type == "application/pdf":
+        context = extract_text_from_pdf(uploaded_file)
+    elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+        context = extract_text_from_docx(uploaded_file)
+    else:
+        # TXT fájl feltöltése esetén
+        try:
+            context = uploaded_file.read().decode("utf-8")
+        except UnicodeDecodeError:
+            context = uploaded_file.read().decode("ISO-8859-1")
+
     # Teljes dokumentum megjelenítése
     st.write("### Feltöltött dokumentum tartalma:")
     st.write(context)