Update app.py

app.py

@@ -17,11 +17,14 @@ def segment_text(input_text, multi_doc_input):
         results["input_text"] = {"segments": sentences}
     
     if multi_doc_input:
-        # Process multiple documents
-        documents = [doc.strip() for doc in multi_doc_input.split('\n') if doc.strip()]
-        for i, doc in enumerate(documents, 1):
-            sentences = sat.split(doc)
-            results[f"row_{i}"] = {"segments": sentences}
+        try:
+            # Parse the JSON input
+            documents = json.loads(multi_doc_input)
+            for key, doc in documents.items():
+                sentences = sat.split(doc)
+                results[f"doc_{key}"] = {"segments": sentences}
+        except json.JSONDecodeError:
+            results["error"] = "Invalid JSON input for multiple documents"
     
     # Create a JSON object with the results
     json_output = json.dumps(results, indent=2)
@@ -33,16 +36,16 @@ iface = gr.Interface(
     fn=segment_text,
     inputs=[
         gr.Textbox(lines=5, label="Input Text (Optional)"),
-        gr.Textbox(lines=10, label="Multiple Documents (Optional, one per line)")
+        gr.Textbox(lines=10, label="Multiple Documents JSON (Optional)", placeholder='{"1": "Document 1 text", "2": "Document 2 text"}')
     ],
     outputs=gr.JSON(label="Segmented Text (JSON)"),
     title="Text Segmentation with SaT",
-    description="This app uses the SaT (Segment any Text) model to split input text into sentences and return the result as JSON. You can input text directly or provide multiple documents (one per line). All credits to the respective author(s). Github: https://github.com/segment-any-text/wtpsplit/tree/main",
+    description="This app uses the SaT (Segment any Text) model to split input text into sentences and return the result as JSON. You can input text directly or provide multiple documents as JSON. All credits to the respective author(s). Github: https://github.com/segment-any-text/wtpsplit/tree/main",
     examples=[
         ["This is a test This is another test.", ""],
         ["Hello this is a test But this is different now Now the next one starts looool", ""],
         ["The quick brown fox jumps over the lazy dog It was the best of times, it was the worst of times", ""],
-        ["", "Document 1 first sentence Document 1 second sentence\nDocument 2 only sentence\nDocument 3 first Document 3 second"]
+        ["", '{"1": "Document 1 first sentence Document 1 second sentence", "2": "Document 2 only sentence", "3": "Document 3 first Document 3 second"}']
     ]
 )