Update app.py

app.py

@@ -6,66 +6,41 @@ import textstat
 
 @st.cache_resource
 def load_model():
+    # Load to CPU explicitly
     base = BartForConditionalGeneration.from_pretrained(
         "facebook/bart-large-cnn",
-        torch_dtype=torch.float16,
-        device_map="auto"
+        torch_dtype=torch.float32,
+        device_map=None  # Don't use auto
     )
+    
     model = PeftModel.from_pretrained(base, "ckharche/legaleaze-bart-121k")
     tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+    
+    model.to("cpu")
     model.eval()
     return tokenizer, model
 
 def simplify(text, tokenizer, model):
     prompt = f"simplify: {text}"
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+    
     with torch.inference_mode():
-        outputs = model.generate(**inputs, max_new_tokens=256, num_beams=4, early_stopping=True, length_penalty=0.75)
+        outputs = model.generate(**inputs, max_new_tokens=256, num_beams=4, early_stopping=True)
+    
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# UI
-st.set_page_config(page_title="Legaleaze", layout="wide")
-st.title("Legaleaze: Legal Text Simplifier")
-st.caption("BART-Large + LoRA | Trained on 121k steps (53k asylum cases)")
+# Rest of your UI code...
+st.title("âš–ï¸ Legaleaze")
+tokenizer, model = load_model()
 
-try:
-    tokenizer, model = load_model()
-    
-    col1, col2 = st.columns(2)
-    
-    with col1:
-        st.subheader("Complex Legal Text")
-        legal_text = st.text_area("", height=300, placeholder="Paste legal text here...", key="input")
-        simplify_btn = st.button("Simplify", type="primary", use_container_width=True)
-    
-    with col2:
-        st.subheader("Simplified Output")
-        if simplify_btn and legal_text.strip():
-            with st.spinner("Simplifying..."):
-                simplified = simplify(legal_text, tokenizer, model)
-                st.text_area("", value=simplified, height=300, disabled=True, key="output")
-                
-                # Metrics
-                st.divider()
-                m1, m2, m3 = st.columns(3)
-                orig_fkgl = textstat.flesch_kincaid_grade(legal_text)
-                simp_fkgl = textstat.flesch_kincaid_grade(simplified)
-                improvement = ((orig_fkgl - simp_fkgl) / orig_fkgl) * 100
-                
-                m1.metric("Original Grade Level", f"{orig_fkgl:.1f}")
-                m2.metric("Simplified Grade Level", f"{simp_fkgl:.1f}")
-                m3.metric("Readability ↑", f"{improvement:.0f}%", delta=f"-{orig_fkgl - simp_fkgl:.1f} grades")
-        else:
-            st.info("Output appears here")
-    
-    with st.expander("ℹ️ Model Details"):
-        st.markdown("""
-        - **Architecture**: BART-Large-CNN (406M params) + LoRA (16M trainable)
-        - **Training**: 121k steps on H100/H200 GPUs (Northeastern HPC)
-        - **Dataset**: 53k Canadian asylum case documents
-        - **Performance**: FKGL ↓35% | BERTScore 0.89 | ROUGE-L 0.48
-        """)
+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
 
-except Exception as e:
-    st.error(f"Model loading failed: {e}")
-    st.info("This Space requires GPU runtime. Contact admin if issue persists.")
+with col2:
+    if 'result' in st.session_state:
+        st.text_area("Simplified", st.session_state['result'], height=300)