Update app.py

app.py

@@ -4,6 +4,9 @@ import seaborn as sns
 import gradio as gr
 import requests
 from bs4 import BeautifulSoup
+import io
+import os
+import base64
 
 # Input data with links to Hugging Face repositories
 data_full = [
@@ -52,8 +55,14 @@ def plot_average_scores():
     plt.gca().invert_yaxis()
     plt.grid(axis='x', linestyle='--', alpha=0.7)
     plt.tight_layout()
-    plt.savefig("average_performance.png")
-    return "average_performance.png"
+    
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format='png')
+    img_buffer.seek(0)
+    img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
+    plt.close()
+    return img_base64, "average_performance.png"
+
 
 def plot_task_performance():
     df_full_melted = df_full.melt(id_vars=["Model Configuration", "Model Link"], var_name="Task", value_name="Score")
@@ -70,8 +79,13 @@ def plot_task_performance():
     plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=9)
     plt.grid(axis='y', linestyle='--', alpha=0.7)
     plt.tight_layout()
-    plt.savefig("task_performance.png")
-    return "task_performance.png"
+    
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format='png')
+    img_buffer.seek(0)
+    img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
+    plt.close()
+    return img_base64, "task_performance.png"
 
 def plot_task_specific_top_models():
     top_models = df_full.iloc[:, 2:].idxmax()
@@ -86,8 +100,13 @@ def plot_task_specific_top_models():
     plt.ylabel("Score", fontsize=14)
     plt.grid(axis="y", linestyle="--", alpha=0.7)
     plt.tight_layout()
-    plt.savefig("task_specific_top_models.png")
-    return "task_specific_top_models.png"
+
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format='png')
+    img_buffer.seek(0)
+    img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
+    plt.close()
+    return img_base64, "task_specific_top_models.png"
 
 def scrape_mergekit_config(model_name):
     """
@@ -109,8 +128,62 @@ def plot_heatmap():
     sns.heatmap(df_full.iloc[:, 2:], annot=True, cmap="YlGnBu", xticklabels=columns[2:], yticklabels=df_full["Model Configuration"])
     plt.title("Performance Heatmap", fontsize=16)
     plt.tight_layout()
-    plt.savefig("performance_heatmap.png")
-    return "performance_heatmap.png"
+    
+    img_buffer = io.BytesIO()
+    plt.savefig(img_buffer, format='png')
+    img_buffer.seek(0)
+    img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
+    plt.close()
+    return img_base64, "performance_heatmap.png"
+
+
+def download_yaml(yaml_content, model_name):
+    """
+    Generates a downloadable link for the scraped YAML content.
+    """
+    if "No YAML configuration found" in yaml_content or "Failed to fetch model page" in yaml_content:
+        return None  # Do not return a link if there's no config or a fetch error
+    
+    filename = f"{model_name.replace('/', '_')}_config.yaml"
+    return gr.File(value=yaml_content.encode(), filename=filename)
+
+def download_all_data():
+    # Prepare data to download
+    csv_buffer = io.StringIO()
+    df_full.to_csv(csv_buffer, index=False)
+    csv_data = csv_buffer.getvalue().encode('utf-8')
+    
+    # Prepare all plots
+    average_plot_b64, average_plot_name = plot_average_scores()
+    task_plot_b64, task_plot_name = plot_task_performance()
+    top_models_plot_b64, top_models_plot_name = plot_task_specific_top_models()
+    heatmap_plot_b64, heatmap_plot_name = plot_heatmap()
+
+    plot_dict = {
+        "average_performance": (average_plot_b64, average_plot_name),
+        "task_performance": (task_plot_b64, task_plot_name),
+        "top_models": (top_models_plot_b64, top_models_plot_name),
+        "heatmap": (heatmap_plot_b64, heatmap_plot_name)
+    }
+
+    zip_buffer = io.BytesIO()
+    import zipfile
+    with zipfile.ZipFile(zip_buffer, 'w') as zf:
+        zf.writestr("model_scores.csv", csv_data)
+
+        for name, (b64, filename) in plot_dict.items():
+            img_data = base64.b64decode(b64)
+            zf.writestr(filename, img_data)
+
+        for model_name in df_full["Model Configuration"].to_list():
+            yaml_content = scrape_mergekit_config(model_name)
+            if "No YAML configuration found" not in yaml_content and "Failed to fetch model page" not in yaml_content:
+               zf.writestr(f"{model_name.replace('/', '_')}_config.yaml", yaml_content.encode())
+
+    zip_buffer.seek(0)
+    
+    return zip_buffer, "analysis_data.zip"
+
 
 # Gradio app
 with gr.Blocks() as demo:
@@ -118,28 +191,44 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         btn1 = gr.Button("Show Average Performance")
-        img1 = gr.Image(type="filepath")
-        btn1.click(plot_average_scores, outputs=img1)
-
+        img1 = gr.Image(type="bytes", label="Average Performance Plot")
+        img1_download = gr.File(label="Download Average Performance")
+        btn1.click(plot_average_scores, outputs=[img1,img1_download])
+        
     with gr.Row():
         btn2 = gr.Button("Show Task Performance")
-        img2 = gr.Image(type="filepath")
-        btn2.click(plot_task_performance, outputs=img2)
+        img2 = gr.Image(type="bytes", label="Task Performance Plot")
+        img2_download = gr.File(label="Download Task Performance")
+        btn2.click(plot_task_performance, outputs=[img2, img2_download])
 
     with gr.Row():
         btn3 = gr.Button("Task-Specific Top Models")
-        img3 = gr.Image(type="filepath")
-        btn3.click(plot_task_specific_top_models, outputs=img3)
-
+        img3 = gr.Image(type="bytes", label="Task-Specific Top Models Plot")
+        img3_download = gr.File(label="Download Top Models")
+        btn3.click(plot_task_specific_top_models, outputs=[img3, img3_download])
+    
     with gr.Row():
-        btn4 = gr.Button("Plot Performance Heatmap")
-        heatmap_img = gr.Image(type="filepath")
-        btn4.click(plot_heatmap, outputs=heatmap_img)
+      btn4 = gr.Button("Plot Performance Heatmap")
+      heatmap_img = gr.Image(type="bytes", label="Performance Heatmap")
+      heatmap_download = gr.File(label="Download Heatmap")
+      btn4.click(plot_heatmap, outputs=[heatmap_img, heatmap_download])
 
     with gr.Row():
         model_selector = gr.Dropdown(choices=df_full["Model Configuration"].tolist(), label="Select a Model")
-        scrape_btn = gr.Button("Scrape MergeKit Configuration")
-        yaml_output = gr.Textbox(lines=10, placeholder="YAML Configuration will appear here.")
-        scrape_btn.click(scrape_mergekit_config, inputs=model_selector, outputs=yaml_output)
+        with gr.Column():
+          scrape_btn = gr.Button("Scrape MergeKit Configuration")
+          yaml_output = gr.Textbox(lines=10, placeholder="YAML Configuration will appear here.")
+          scrape_btn.click(scrape_mergekit_config, inputs=model_selector, outputs=yaml_output)
+        with gr.Column():
+          save_yaml_btn = gr.Button("Save MergeKit Configuration")
+          yaml_download = gr.File(label="Download MergeKit Configuration")
+          save_yaml_btn.click(download_yaml, inputs=[yaml_output, model_selector], outputs=yaml_download)
+
+
+    with gr.Row():
+        download_all_btn = gr.Button("Download Everything")
+        all_downloads = gr.File(label="Download All Data")
+        download_all_btn.click(download_all_data, outputs=all_downloads)
+
 
-demo.launch()
+demo.launch()