final

app.py

@@ -190,11 +190,52 @@ def save_model(
 
     # Format it for display in markdown with proper styling
     model_architecture_info = f"""
-    <div class="model-architecture" style="max-height: 500px; overflow-y: auto; overflow-x: auto; background-color: #f5f5f5; padding: 5px; border-radius: 8px; font-family: monospace; white-space: pre-wrap;">
+    <div class="model-architecture-container" style="margin-top: 20px; margin-bottom: 20px; background-color: #f8f9fa; padding: 15px; border-radius: 8px; border-left: 4px solid #4CAF50;">
+        <h3 style="margin-top: 0; color: #2E7D32;">📋 Model Architecture</h3>
+        <div class="model-architecture" style="max-height: 500px; overflow-y: auto; overflow-x: auto; background-color: #f5f5f5; padding: 5px; border-radius: 8px; font-family: monospace; white-space: pre-wrap;">
         <div style="line-height: 1.2; font-size: 0.75em;">{model_architecture_str_html}</div>
+        </div>
+    </div>
+    """
+
+    code = f"""
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model = AutoModelForCausalLM.from_pretrained(
+    "{repo_name}",
+    device_map="auto"
+)
+
+tokenizer = AutoTokenizer.from_pretrained("{model_name}")
+
+inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)
+outputs = model.generate(inputs.input_ids, max_length=50)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+"""
+    # First replace newlines in the code
+    code_html = code.replace("\n", "<br/>")
+    
+    usage_example = f"""
+    <div class="usage-example" style="margin-top: 20px; margin-bottom: 20px; background-color: #f8f9fa; padding: 15px; border-radius: 8px; border-left: 4px solid #4CAF50;">
+        <h3 style="margin-top: 0; color: #2E7D32;">📝 How to use this model</h3>
+        <div class="model-architecture" style="max-height: 500px; overflow-y: auto; overflow-x: auto; background-color: #f5f5f5; padding: 5px; border-radius: 8px; font-family: monospace; font-size: 0.75em; white-space: pre-wrap;">{code_html}</div>
+    </div>
+    """
+    model_size_info = f"""
+    <div class="model-size-info" style="margin-top: 20px; margin-bottom: 20px; background-color: #f8f9fa; padding: 15px; border-radius: 8px; border-left: 4px solid #4CAF50;">
+        <h3 style="margin-top: 0; color: #2E7D32;">📦 Model Size</h3>
+        <p>Original (bf16)≈ {original_size_gb} GB → Quantized ≈ {get_model_size(model)} GB</p>
     </div>
     """
-    return f'🎉 Quantized Model <br/><h1> 🚀 DONE</h1><br/>Find your repo here: <a href="https://huggingface.co/{repo_name}" target="_blank" style="text-decoration:underline">{repo_name}</a><br/><br/>📋 Model Architecture<br/>{model_architecture_info}<br/><br/>📦 Model Size <br/><br/> Original (bf16)≈ {original_size_gb} GB → Quantized ≈ {get_model_size(model)} GB'
+
+    repo_link = f"""
+    <div class="repo-link" style="margin-top: 20px; margin-bottom: 20px; background-color: #f8f9fa; padding: 15px; border-radius: 8px; border-left: 4px solid #4CAF50;">
+        <h3 style="margin-top: 0; color: #2E7D32;">🔗 Repository Link</h3>
+        <p>Find your repo here: <a href="https://huggingface.co/{repo_name}" target="_blank" style="text-decoration:underline">{repo_name}</a></p>
+    </div>
+    """
+    return f'<h1>🎉 Quantization Completed</h1><br/>{repo_link}{model_architecture_info}{model_size_info}{usage_example}'
 
 
 def quantize_and_save(
@@ -508,7 +549,8 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
         # 🤗 BitsAndBytes Quantizer : Create your own BNB Quants ! ✨
 
 
-        
+        <br/>
+        <br/>
         """
     )
 
@@ -597,7 +639,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
                 "🚀 Quantize and Push to the Hub", variant="primary"
             )
             output_link = gr.Markdown(
-                "🔗 Quantized Model", container=True, min_height=100
+                "🔗 Quantized Model Info", container=True, min_height=200
             )
 
     quantize_button.click(