Update app.py

app.py

@@ -1,14 +1,15 @@
 import streamlit as st
-import torch
 from huggingface_hub import hf_hub_download
 from utils import *
 
+# Download the models
 HengamTransW = hf_hub_download(repo_id="kargaranamir/Hengam", filename="HengamTransW.pth")
 HengamTransA = hf_hub_download(repo_id="kargaranamir/Hengam", filename="HengamTransA.pth")
 
-ner_W = NER(model_path=HengamTransW, tags=['B-TIM', 'I-TIM', 'B-DAT', 'I-DAT', 'O'])
-ner_A = NER(model_path=HengamTransA, tags=['B-TIM', 'I-TIM', 'B-DAT', 'I-DAT', 'O'])
-
+# Load models lazily
+@st.cache(allow_output_mutation=True)
+def load_ner_model(model_path):
+    return NER(model_path=model_path, tags=['B-TIM', 'I-TIM', 'B-DAT', 'I-DAT', 'O'])
 
 # APP
 st.set_page_config(
@@ -16,71 +17,63 @@ st.set_page_config(
     page_icon="🕒",
 )
 
+# Layout adjustments
+st.markdown(
+    """
+<style>
+.reportview-container .main .block-container {
+    max-width: 1400px;
+}
+</style>
+""",
+    unsafe_allow_html=True,
+)
 
-def _max_width_():
-    max_width_str = f"max-width: 1400px;"
-    st.markdown(
-        f"""
-    <style>
-    .reportview-container .main .block-container{{
-        {max_width_str}
-    }}
-    </style>    
-    """,
-        unsafe_allow_html=True,
-    )
-
-
-_max_width_()
-
-c30, c31, c32 = st.columns([2.5, 1, 3])
-
-with c30:
-    st.title("🕒 Hengam")
-    st.header("")
-
+st.markdown("## **🕒 Hengam**")
+st.write("")
 with st.expander("ℹ️ - About this app", expanded=True):
     st.write(
         """     
--   Online Demo for Hengam: An Adversarially Trained Transformer for Persian Temporal Tagging [Code](https://github.com/kargaranamir/hengam)!
--   This paper introduces Hengam, an adversarially trained transformer for Persian temporal tagging outperforming state-of-the-art approaches on a diverse and manually created dataset.
-	    """
+- Online Demo for Hengam: An Adversarially Trained Transformer for Persian Temporal Tagging [Code](https://github.com/kargaranamir/hengam)!
+- This paper introduces Hengam, an adversarially trained transformer for Persian temporal tagging outperforming state-of-the-art approaches on a diverse and manually created dataset.
+    """
     )
     st.markdown("")
 
 st.markdown("")
-st.markdown("## **📌 Paste any persian (farsi) text you want to extract its temporal markers.**")
+st.markdown("## **📌 Paste any Persian (Farsi) text you want to extract its temporal markers.**")
 with st.form(key="my_form"):
-    ce, c1, ce, c2, c3 = st.columns([0.05, 1.5, 0.05, 4, 0.05])
+    c1, c2 = st.columns([2, 3])
+    
     with c1:
-        ModelType = st.radio(
+        st.write("Please wait while loading the model...")
+        model_paths = {
+            "HengamTransW.pth": "HengamTransW.pth",
+            "HengamTransA.pth": "HengamTransA.pth",
+        }
+        
+        default_model = "HengamTransA.pth"  # Set the default model
+        ModelType = st.selectbox(
             "Choose your model",
-            ["HengamTransW.pth", "HengamTransA.pth"],
-            help="At present, you can choose between 2 models (HengamTransW or HengamTransA) to exrtact temporal markers. More to come!",
+            list(model_paths.keys()),
+            index=list(model_paths.keys()).index(default_model),  # Select the default model
+            help="At present, you can choose between 2 models (HengamTransW or HengamTransA) to extract temporal markers. More to come!",
         )
-
-        if ModelType == 'HengamTransW.pth':
-            ner = ner_W
-        else:
-            ner = ner_A
+        
+        ner = load_ner_model(model_paths[ModelType])
+        st.empty()  # Clear the "Please wait" message
 
     with c2:
         doc = st.text_area(
             "Paste your text below",
-            '''ساعت ۸ صبح من و علی قرار گذاشتیم که شانزده بهمن ۱۳۷۵ هم دیگر را در دوشنبه بازار ببینیم.''',
-            height=110,
+            "Example: ساعت ۸ صبح من و علی قرار گذاشتیم که به دوشنبه بازار بریم ...",
+            height=80,
         )
 
         submit_button = st.form_submit_button(label="✨ Extract Temporal Markers!")
 
-if not submit_button:
-    st.stop()
-
-result = "\n".join([str(t) for t in ner(doc)])
-
-st.markdown("## **🎈Check results**")
-
-c1, c2, c3 = st.columns([1, 3, 1])
-
-with c2:
-    st.write(result)
+if submit_button:
+    result = extract_temporal_markers(doc, ner)
+    st.write("")  # Add vertical spacing
+    st.markdown("## **🎈Check results**")
+    st.code(result, height=300)  # Adjust the height here