Update app.py

app.py

@@ -1,48 +1,144 @@
-import os
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import HfApi
+import gc
+import sys
 
-def prune_to_single_layer_and_push(base_model_id: str, new_repo_id: str, make_private: bool):
-    hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        return "❌ Set HF_TOKEN secret in Space Settings first!"
-    
-    if not base_model_id or not new_repo_id:
-        return "❌ Fill model and repo name"
+def get_model_size_mb(model):
+    """Rough estimate of model size in MB (parameters only)"""
+    param_size = 0
+    for param in model.parameters():
+        param_size += param.nelement() * param.element_size()
+    return round(param_size / (1024 ** 2), 1)
+
+def prune_to_single_layer(model_id: str):
+    status_lines = []
+    status_lines.append(f"Loading base model: {model_id}")
     
     try:
-        # Load safely on CPU (free Space friendly)
+        # Load on CPU with low memory usage flags
         model = AutoModelForCausalLM.from_pretrained(
-            base_model_id,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            model_id,
+            torch_dtype=torch.float32,           # float32 = most compatible on CPU
             device_map="cpu",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
+            trust_remote_code=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            trust_remote_code=True
         )
-        tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)
         
-        # === SINGLE LAYER MAGIC ===
-        if hasattr(model, "model") and hasattr(model.model, "layers"):
-            model.model.layers = torch.nn.ModuleList([model.model.layers[-1]])  # keep LAST layer only
-            model.config.num_hidden_layers = 1
-        else:
-            return "❌ Model type not supported (needs standard .model.layers)"
+        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
         
-        # Push
-        model.push_to_hub(new_repo_id, private=make_private, token=hf_token, safe_serialization=True)
-        tokenizer.push_to_hub(new_repo_id, private=make_private, token=hf_token)
+        # Optional: clear intermediate tensors if possible
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         
-        link = f"https://huggingface.co/{new_repo_id}"
-        return f"""✅ SUCCESS!
+        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)}%")
+        
+        # Quick generation smoke test
+        try:
+            inputs = tokenizer("Hello, the future of single-layer models is", return_tensors="pt")
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs.to(model.device),
+                    max_new_tokens=40,
+                    do_sample=False,
+                    temperature=0.0
+                )
+            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():
+            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:
+            del model
+            del tokenizer
+            gc.collect()
+        except:
+            pass
+
+# ────────────────────────────────────────────────
+#                 Gradio Interface
+# ────────────────────────────────────────────────
 
-Single-layer model is live at:
-**{link}**
+CSS = """
+.gradio-container { max-width: 780px !important; }
+"""
 
-Size: ~0.3–0.8 GB → runs at 40–100+ tokens/s on any CPU!
+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",
+            "Qwen/Qwen2.5-1.5B-Instruct",
+            "meta-llama/Llama-3.2-1B-Instruct",
+            "google/gemma-2-2b-it",
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        ],
+        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.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
+    """)
 
-Load it anywhere:
-```python
-from transformers import AutoModelForCausalLM
-model = AutoModelForCausalLM.from_pretrained("{new_repo_id}")
+demo.launch()