Update app.py

app.py

@@ -1,45 +1,18 @@
 import streamlit as st
+import pdfplumber
+import docx
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 st.title("Jogi Dokumentumok AI Válaszgenerálás")
 
-# Egyedi session azonosító létrehozása
-session_id = st.session_state.session_id if "session_id" in st.session_state else 0
-
 # Előző beszélgetések tárolása session state-ben
 if "previous_conversations" not in st.session_state:
     st.session_state.previous_conversations = []
 
 # Korábbi beszélgetések törlése
-if st.button("Korábbi Beszélgetések Törlése"):
+if st.checkbox("Korábbi Beszélgetések Törlése"):
     st.session_state.previous_conversations = []
 
-# Korábbi beszélgetések mentése és betöltése
-if st.checkbox("Korábbi Beszélgetések Mentése"):
-    with open(f"previous_conversations_{session_id}.txt", "w") as file:
-        for conv in st.session_state.previous_conversations:
-            file.write(f"{conv}\n")
-    st.write("A korábbi beszélgetések sikeresen el lettek mentve.")
-
-if st.checkbox("Korábbi Beszélgetések Betöltése"):
-    previous_conversations_file = st.file_uploader("Válassz ki egy fájlt a korábbi beszélgetések betöltéséhez")
-    if previous_conversations_file is not None:
-        with previous_conversations_file as file:
-            for line in file:
-                st.session_state.previous_conversations.append(line.strip())
-        st.write("A korábbi beszélgetések sikeresen betöltve.")
-
-# Korábbi beszélgetések keresése
-search_term = st.text_input("Keresés a Korábbi Beszélgetések Között", "")
-if st.button("Keresés"):
-    search_results = [conv for conv in st.session_state.previous_conversations if search_term in conv]
-    if search_results:
-        st.subheader("Keresési Találatok:")
-        for result in search_results:
-            st.write(result)
-    else:
-        st.write("Nincs találat.")
-
 # Tokenizáló modell és válaszgeneráló modell választása
 tokenizer_name = st.selectbox("Tokenizáló Modell", ["bert-base-uncased", "bert-base-multilingual-cased"])
 generator_name = st.selectbox("Válaszgeneráló Modell", ["microsoft/phi-2", "gpt2"])
@@ -47,25 +20,46 @@ generator_name = st.selectbox("Válaszgeneráló Modell", ["microsoft/phi-2", "g
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
 model = AutoModelForCausalLM.from_pretrained(generator_name, trust_remote_code=True)
 
-# Input mező a jogi dokumentumhoz
-document_text = st.text_area("Jogi Dokumentum", "")
-
-# Előző beszélgetésekhez csatolható kontextus
-context = st.text_area("Korábbi Beszélgetéshez Tartozó Kontextus", "")
+# Dokumentum feltöltése drag and drop segítségével
+document_file = st.file_uploader("Húzd ide a dokumentumot vagy kattints a feltöltéshez", type=["pdf", "docx", "doc"])
 
-# Válaszgenerálás
-if st.button("Generálj Választ"):
-    if context:
-        input_text = f"{context} {document_text}"
+if document_file is not None:
+    if document_file.type == "application/pdf":
+        with pdfplumber.open(document_file) as pdf:
+            document_text = ""
+            for page in pdf.pages:
+                document_text += page.extract_text()
+    elif document_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+        docx_file = docx.Document(document_file)
+        document_text = ""
+        for paragraph in docx_file.paragraphs:
+            document_text += paragraph.text
+    elif document_file.type == "application/msword":
+        doc_file = docx.Document(document_file)
+        document_text = ""
+        for paragraph in doc_file.paragraphs:
+            document_text += paragraph.text
     else:
-        input_text = document_text
+        st.error("A fájltípus nem támogatott. Kérlek válassz ki egy PDF, DOCX vagy DOC fájlt!")
+
+    # Előző beszélgetésekhez csatolható kontextus
+    context = st.text_area("Korábbi Beszélgetéshez Tartozó Kontextus", "")
+
+    # Válaszgenerálás
+    if st.button("Generálj Választ"):
+        if context:
+            input_text = f"{context} {document_text}"
+        else:
+            input_text = document_text
 
-    # Model használata a válasz generálásához
-    response = generate_response(input_text, tokenizer, model)
+        # Model használata a válasz generálásához
+        response = generate_response(input_text, tokenizer, model)
 
-    # Válasz megjelenítése
-    st.subheader("Generált Válasz:")
-    st.write(response)
+        # Válasz megjelenítése
+        st.subheader("Generált Válasz:")
+        st.write(response)
 
-    # Aktuális beszélgetés hozzáadása az előző beszélgetésekhez
-    st.session_state.previous_conversations.append({"session_id": session_id, "input_text": input_text, "response": response})
+        # Aktuális beszélgetés hozzáadása az előző beszélgetésekhez
+        st.session_state.previous_conversations.append({"input_text": input_text, "response": response})
+else:
+    st.warning("Kérlek válassz ki egy dokumentumfájlt!")