Update app.py

app.py

@@ -75,12 +75,12 @@ def orthogonalize_matrix(matrix, vec, weight):
 # --- Streamlit UI ---
 
 st.title("LLM Auto Abliteration")
-st.markdown("🥠Recommended to run edge-device LLMs (e.g., 1B, 1.5B, 0.5B).")
-st.markdown("🥠And also recommended to duplicate the space for seamless usage!")
-st.markdown("🥠This app allows you to manually input parameters to modify a language model's behavior by abliterating its weights.")
-st.markdown("📍Credits: Thanks to **[Maxime Labonne](https://huggingface.co/mlabonne)**")
+st.markdown("🥠 Recommended for edge-device LLMs (e.g., 1B, 1.5B, 0.5B).")
+st.markdown("🥠 Duplicate the space for seamless usage!")
+st.markdown("🥠 This app allows you to manually input parameters to modify a language model's behavior by abliterating its weights.")
+st.markdown("📍 Credits: Thanks to **[Maxime Labonne](https://huggingface.co/mlabonne)**")
 
-# Debugging window (will update logs during the process)
+# Debugging window to show log messages
 debug_log = []
 debug_placeholder = st.empty()
 def update_debug(msg):
@@ -132,14 +132,20 @@ if st.button("Run Abliteration"):
     
     with st.spinner("Loading model..."):
         try:
+            # Pass HF token if provided to load private or restricted models
             model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
                 device_map="auto",
                 torch_dtype=torch_dtype,
                 attn_implementation=attn_implementation,
-                trust_remote_code=True
+                trust_remote_code=True,
+                use_auth_token=hf_token if hf_token else None
             ).eval()
-            tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+            tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID,
+                trust_remote_code=True,
+                use_auth_token=hf_token if hf_token else None
+            )
             update_debug("Model and tokenizer loaded successfully.")
         except Exception as e:
             update_debug(f"Error loading model: {e}")
@@ -252,17 +258,16 @@ if st.button("Run Abliteration"):
     st.text_area("Response after abliteration:", response_after, height=150)
     update_debug("Generated response after abliteration.")
     
-    st.write("### (Optional) Pushing Model to Hugging Face Hub")
-    if st.checkbox("Push model to HF Hub?"):
-        try:
-            model_name = MODEL_ID.split("/")[-1] + "-abliterated"
-            model.push_to_hub(model_name, private=PRIVATE_UPLOAD)
-            tokenizer.push_to_hub(model_name, private=PRIVATE_UPLOAD)
-            st.success(f"Model pushed as {model_name}")
-            update_debug(f"Model pushed to HF Hub as {model_name}.")
-        except Exception as e:
-            st.error(f"Error while pushing model: {e}")
-            update_debug(f"Error while pushing model: {e}")
+    st.write("### Pushing Model to Hugging Face Hub")
+    try:
+        model_name = MODEL_ID.split("/")[-1] + "-abliterated"
+        model.push_to_hub(model_name, private=PRIVATE_UPLOAD)
+        tokenizer.push_to_hub(model_name, private=PRIVATE_UPLOAD)
+        st.success(f"Model automatically pushed as {model_name}")
+        update_debug(f"Model automatically pushed to HF Hub as {model_name}.")
+    except Exception as e:
+        st.error(f"Error while pushing model: {e}")
+        update_debug(f"Error while pushing model: {e}")
     
     st.success("Abliteration process complete!")
     update_debug("Abliteration process complete.")