Always start off updated doc with original doc

app.py

@@ -100,9 +100,17 @@ def highlight_edits():
     st.write("Prompt:", prompt)
     cols = st.columns(2)
     with cols[0]:
-        doc = st.text_area("Document", "Deep learning neural network technology advances are pretty cool if you are careful to use it in ways that don't take stuff from people.", height=300)
+        def on_doc_change():
+            """Set the updated_doc to start with the original doc."""
+            st.session_state['updated_doc'] = st.session_state['doc']
+        doc = st.text_area(
+            "Document",
+            "Deep learning neural network technology advances are pretty cool if you are careful to use it in ways that don't take stuff from people.",
+            height=300, key='doc',
+            on_change=on_doc_change
+        )
     with cols[1]:
-        updated_doc = st.text_area("Updated Doc", placeholder="Your edited document. Leave this blank to use your original document.", height=300)
+        updated_doc = st.text_area("Updated Doc", placeholder="Your edited document. Leave this blank to use your original document.", height=300, key='updated_doc')
 
     spans = get_highlights(prompt, doc, updated_doc)