Update app.py

app.py

@@ -25,6 +25,9 @@ document_text = st.text_area("Jogi Dokumentum", "")
 # Dokumentum típusának megadása
 document_type = st.text_input("Dokumentum Típusa", "")
 
+# Szövegszegmentáció
+segmentation_type = st.selectbox("Szövegszegmentáció", ["Címek", "Bekezdések", "Fejezetek"])
+
 # Generálási paraméterek testreszabása
 max_length = st.slider("Maximális Válasz Hossz", min_value=10, max_value=200, value=50, step=10)
 num_return_sequences = st.slider("Válaszok Száma", min_value=1, max_value=5, value=3, step=1)
@@ -38,29 +41,42 @@ if st.button("Tokenizálás és Válaszgenerálás"):
         file.write(str(tokenized_document))
     st.write("A tokenizált dokumentum sikeresen el lett mentve.")
 
+    # Szövegszegmentáció
+    if segmentation_type == "Címek":
+        segments = extract_titles(document_text)
+    elif segmentation_type == "Bekezdések":
+        segments = extract_paragraphs(document_text)
+    elif segmentation_type == "Fejezetek":
+        segments = extract_chapters(document_text)
+
     # Válasz generálása
-    input_ids = tokenized_document["input_ids"]
-    output = model.generate(input_ids, max_length=max_length, num_return_sequences=num_return_sequences, temperature=temperature)
-    st.subheader("Generált Válaszok:")
-    for i, generated_sequence in enumerate(output):
-        response_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
-        st.write(f"Válasz {i+1}: {response_text}")
-        if hasattr(model, "get_output_embeddings"):
-            logits = model.get_output_embeddings()(generated_sequence)
-            probs = logits.softmax(dim=-1)
-            st.write("Valószínűségek:")
-            st.table(probs.tolist())
-        if st.button(f"Mentés Válasz {i+1}"):
-            file_name = st.text_input("Fájlnév:", f"generated_response_{i+1}.txt")
-            with open(file_name, "w", encoding="utf-8") as file:
-                file.write(response_text)
-            st.write(f"A válasz {i+1} sikeresen el lett mentve fájlba.")
+    for segment in segments:
+        input_ids = tokenizer.encode(segment, return_tensors="pt")
+        output = model.generate(input_ids, max_length=max_length, num_return_sequences=num_return_sequences, temperature=temperature)
+        st.subheader(f"{segment} - Generált Válaszok:")
+        for i, generated_sequence in enumerate(output):
+            response_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
+            st.write(f"Válasz {i+1}: {response_text}")
+            if hasattr(model, "get_output_embeddings"):
+                logits = model.get_output_embeddings()(generated_sequence)
+                probs = logits.softmax(dim=-1)
+                st.write("Valószínűségek:")
+                st.table(probs.tolist())
+            if st.button(f"Mentés Válasz {i+1}"):
+                file_name = st.text_input("Fájlnév:", f"generated_response_{i+1}.txt")
+                with open(file_name, "w", encoding="utf-8") as file:
+                    file.write(response_text)
+                st.write(f"A válasz {i+1} sikeresen el lett mentve fájlba.")
 
     # Generált válaszok egyetlen fájlba mentése
     if st.button("Mentés Összes Válasz"):
         with open("all_generated_responses.txt", "w", encoding="utf-8") as file:
-            for i, generated_sequence in enumerate(output):
-                generated_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
-                file.write(f"Válasz {i+1}:\n")
-                file.write(generated_text + "\n\n")
+            for segment in segments:
+                input_ids = tokenizer.encode(segment, return_tensors="pt")
+                output = model.generate(input_ids, max_length=max_length, num_return_sequences=num_return_sequences, temperature=temperature)
+                file.write(f"{segment}:\n")
+                for i, generated_sequence in enumerate(output):
+                    generated_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
+                    file.write(f"Válasz {i+1}:\n")
+                    file.write(generated_text + "\n\n")
         st.write("Az összes válasz sikeresen el lett mentve egy fájlba.")