Update app.py

app.py

@@ -6,14 +6,14 @@ import textstat
 
 @st.cache_resource
 def load_model():
-    # Load to CPU explicitly
     base = BartForConditionalGeneration.from_pretrained(
         "facebook/bart-large-cnn",
         torch_dtype=torch.float32,
-        device_map=None  # Don't use auto
+        device_map=None
     )
     
-    model = PeftModel.from_pretrained(base, "ckharche/Legaleaze")
+    # Load from local checkpoint folder
+    model = PeftModel.from_pretrained(base, "./checkpoint")
     tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
     
     model.to("cpu")
@@ -29,18 +29,36 @@ def simplify(text, tokenizer, model):
     
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Rest of your UI code...
-st.title("Legaleaze")
-tokenizer, model = load_model()
+st.set_page_config(page_title="Legaleaze", layout="wide")
+st.title("âš–ï¸ Legaleaze: Legal Text Simplifier")
+st.caption("BART-Large + LoRA | 121k steps on asylum cases")
 
-col1, col2 = st.columns(2)
-with col1:
-    text = st.text_area("Complex Legal Text", height=300)
-    if st.button("Simplify"):
-        with st.spinner("Processing (20-30s on CPU)..."):
-            result = simplify(text, tokenizer, model)
-            st.session_state['result'] = result
+try:
+    tokenizer, model = load_model()
+    
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.subheader("Complex Legal Text")
+        text = st.text_area("", height=300, placeholder="Paste legal text here...")
+        btn = st.button("Simplify", type="primary", use_container_width=True)
+    
+    with col2:
+        st.subheader("Simplified Output")
+        if btn and text.strip():
+            with st.spinner("Simplifying (30s on CPU)..."):
+                result = simplify(text, tokenizer, model)
+                st.text_area("", value=result, height=300, disabled=True)
+                
+                st.divider()
+                m1, m2, m3 = st.columns(3)
+                orig = textstat.flesch_kincaid_grade(text)
+                simp = textstat.flesch_kincaid_grade(result)
+                m1.metric("Original FKGL", f"{orig:.1f}")
+                m2.metric("Simplified FKGL", f"{simp:.1f}")
+                m3.metric("Improvement", f"{((orig-simp)/orig*100):.0f}%")
+        else:
+            st.info("Simplified text appears here")
 
-with col2:
-    if 'result' in st.session_state:
-        st.text_area("Simplified", st.session_state['result'], height=300)
+except Exception as e:
+    st.error(f"Error: {e}")