add arch

app.py

@@ -113,7 +113,28 @@ def save_model(model, model_name, quant_type_4, double_quant_4, compute_type_4,
             repo_id=repo_name,
             repo_type="model",
         )
-    return f'<h1> 🤗 DONE</h1><br/>Find your repo here: <a href="https://huggingface.co/{repo_name}" target="_blank" style="text-decoration:underline">{repo_name}</a>'
+    # Get model architecture as string
+    import io
+    from contextlib import redirect_stdout
+    import html
+    
+    # Capture the model architecture string
+    f = io.StringIO()
+    with redirect_stdout(f):
+        print(model)
+    model_architecture_str = f.getvalue()
+    
+    # Escape HTML characters and format with line breaks
+    model_architecture_str_html = html.escape(model_architecture_str).replace('\n', '<br/>')
+    
+    # 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 style="line-height: 1.2; font-size: 0.75em;">{model_architecture_str_html}</div>
+    </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}'
 
 def quantize_and_save(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None, model_name, quant_type_4, double_quant_4, compute_type_4, quant_storage_4, quantized_model_name, public):
     if oauth_token is None : 
@@ -129,7 +150,7 @@ def quantize_and_save(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToke
             <h3>❌ Authentication Error</h3>
             <p>Please sign in to your HuggingFace account to use the quantizer.</p>
         </div>
-        """
+        """ 
     exists_message = check_model_exists(oauth_token, profile.username, model_name, quantized_model_name)
     if exists_message : 
         return f"""
@@ -415,8 +436,8 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
 
         with gr.Column():
             quantize_button = gr.Button("🚀 Quantize and Push to the Hub", variant="primary")
-            output_link = gr.Markdown(label="🔗 Quantized Model Link", container=True, min_height=80)   
-    
+            output_link = gr.Markdown("🔗 Quantized Model", container=True, min_height=100)   
+            
     quantize_button.click(
         fn=quantize_and_save,
         inputs=[model_name, quant_type_4, double_quant_4, compute_type_4, quant_storage_4, quantized_model_name, public],