Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gc
-import sys
 
 def get_model_size_mb(model):
     """Rough estimate of model size in MB (parameters only)"""
@@ -15,11 +14,15 @@ def prune_to_single_layer(model_id: str):
     status_lines = []
     status_lines.append(f"Loading base model: {model_id}")
     
+    # Initialize as None to prevent UnboundLocalError in the finally block
+    model = None
+    tokenizer = None
+        
     try:
-        # Load on CPU with low memory usage flags
+        # Load on CPU with bfloat16 to prevent Out of Memory (OOM) on free spaces
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
-            torch_dtype=torch.float32,           # float32 = most compatible on CPU
+            torch_dtype=torch.bfloat16, 
             device_map="cpu",
             low_cpu_mem_usage=True,
             trust_remote_code=True
@@ -28,32 +31,33 @@ def prune_to_single_layer(model_id: str):
             model_id,
             trust_remote_code=True
         )
-        
+                
         orig_layers = len(model.model.layers) if hasattr(model.model, "layers") else "unknown"
         orig_size_mb = get_model_size_mb(model)
         status_lines.append(f"→ Original layers: {orig_layers}")
         status_lines.append(f"→ Original size (approx): {orig_size_mb} MB")
-        
+                
         # ────────────────────────────────────────────────
         # Core pruning step
         if not hasattr(model, "model") or not hasattr(model.model, "layers"):
             return "\n".join(status_lines) + "\n\n❌ Model architecture not supported (no .model.layers found)"
-        
+                
         # Keep only the LAST layer
         model.model.layers = torch.nn.ModuleList([model.model.layers[-1]])
         model.config.num_hidden_layers = 1
-        
-        # Optional: clear intermediate tensors if possible
+                
+        # Clear intermediate tensors
         gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
+                
         new_layers = len(model.model.layers)
         new_size_mb = get_model_size_mb(model)
         status_lines.append(f"→ After pruning: {new_layers} layer")
         status_lines.append(f"→ New size (approx): {new_size_mb} MB")
-        status_lines.append(f"→ Size reduction: ~{round((orig_size_mb - new_size_mb)/orig_size_mb*100)}%")
         
+        if orig_size_mb > 0:
+            reduction = round((orig_size_mb - new_size_mb) / orig_size_mb * 100)
+            status_lines.append(f"→ Size reduction: ~{reduction}%")
+                
         # Quick generation smoke test
         try:
             inputs = tokenizer("Hello, the future of single-layer models is", return_tensors="pt")
@@ -61,52 +65,46 @@ def prune_to_single_layer(model_id: str):
                 outputs = model.generate(
                     **inputs.to(model.device),
                     max_new_tokens=40,
-                    do_sample=False,
-                    temperature=0.0
+                    do_sample=False, # Temperature removed to prevent conflict
+                    pad_token_id=tokenizer.eos_token_id # Prevents warnings/crashes on Llama
                 )
             text = tokenizer.decode(outputs[0], skip_special_tokens=True)
             status_lines.append("\nQuick generation test (should be at least semi-coherent):")
             status_lines.append("→ " + text.strip())
         except Exception as gen_e:
             status_lines.append(f"\nGeneration test failed: {str(gen_e)} (still might be usable)")
-        
+                
         status_lines.append("\nPruning appears successful ✓")
         status_lines.append("You can now safely close this tab or try another model.")
-        
+                
         return "\n".join(status_lines)
-    
+        
     except Exception as e:
         err_msg = str(e)
-        if "out of memory" in err_msg.lower():
+        if "out of memory" in err_msg.lower() or "killed" in err_msg.lower():
             return "\n".join(status_lines) + "\n\n❌ Out of memory — try an even smaller model (0.5B class)"
         return "\n".join(status_lines) + f"\n\n❌ Failed: {err_msg}"
-    
+        
     finally:
-        # Try to free memory even on failure
-        try:
+        # Safely try to free memory even on failure
+        if model is not None:
             del model
+        if tokenizer is not None:
             del tokenizer
-            gc.collect()
-        except:
-            pass
+        gc.collect()
 
 # ────────────────────────────────────────────────
 #                 Gradio Interface
 # ────────────────────────────────────────────────
-
-CSS = """
-.gradio-container { max-width: 780px !important; }
-"""
+CSS = """.gradio-container { max-width: 780px !important; }"""
 
 with gr.Blocks(title="Minimal Single-Layer Pruner", css=CSS, theme=gr.themes.Default()) as demo:
     gr.Markdown("""
     # Single-Layer Pruner (test version)
-    
     Loads a small model → keeps **only the last layer** → shows result + quick generation test.
-    
     **No pushing to Hub yet** — just checking if pruning works reliably.
     """)
-    
+        
     model_choice = gr.Dropdown(
         choices=[
             "Qwen/Qwen2.5-0.5B-Instruct",
@@ -118,27 +116,28 @@ with gr.Blocks(title="Minimal Single-Layer Pruner", css=CSS, theme=gr.themes.Def
         label="Choose small model (0.5B–2B recommended for free CPU Space)",
         value="Qwen/Qwen2.5-0.5B-Instruct"
     )
-    
+        
     status = gr.Textbox(
         label="Pruning log",
         lines=18,
         interactive=False,
         show_copy_button=True
     )
-    
-    btn = gr.Button("Prune to 1 layer → Test", variant="primary", scale=0)
-    
+        
+    btn = gr.Button("Prune to 1 layer → Test", variant="primary")
+        
     btn.click(
         prune_to_single_layer,
         inputs=model_choice,
         outputs=status
     )
-    
+        
     gr.Markdown("""
-    **Tips**  
-    • Start with 0.5B or 1.1B models — they almost always succeed on free Spaces  
-    • The generation test often produces short but semi-sensible text  
-    • Next steps (after this works): add push button, add chat tab, convert to GGUF
+    **Tips**
+      • Start with 0.5B or 1.1B models — they almost always succeed on free Spaces
+      • The generation test often produces short but semi-sensible text
+      • Next steps (after this works): add push button, add chat tab, convert to GGUF
     """)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()