Update app.py

app.py

@@ -7,6 +7,10 @@ from bs4 import BeautifulSoup
 import io
 import os
 import base64
+import zipfile
+from PIL import Image
+from io import BytesIO
+
 
 # Input data with links to Hugging Face repositories
 data_full = [
@@ -61,7 +65,9 @@ def plot_average_scores():
     img_buffer.seek(0)
     img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
     plt.close()
-    return img_base64, "average_performance.png"
+
+    pil_image = Image.open(BytesIO(base64.b64decode(img_base64)))
+    return pil_image, "average_performance.png"
 
 
 def plot_task_performance():
@@ -85,7 +91,8 @@ def plot_task_performance():
     img_buffer.seek(0)
     img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
     plt.close()
-    return img_base64, "task_performance.png"
+    pil_image = Image.open(BytesIO(base64.b64decode(img_base64)))
+    return pil_image, "task_performance.png"
 
 def plot_task_specific_top_models():
     top_models = df_full.iloc[:, 2:].idxmax()
@@ -106,7 +113,8 @@ def plot_task_specific_top_models():
     img_buffer.seek(0)
     img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
     plt.close()
-    return img_base64, "task_specific_top_models.png"
+    pil_image = Image.open(BytesIO(base64.b64decode(img_base64)))
+    return pil_image, "task_specific_top_models.png"
 
 def scrape_mergekit_config(model_name):
     """
@@ -134,7 +142,8 @@ def plot_heatmap():
     img_buffer.seek(0)
     img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
     plt.close()
-    return img_base64, "performance_heatmap.png"
+    pil_image = Image.open(BytesIO(base64.b64decode(img_base64)))
+    return pil_image, "performance_heatmap.png"
 
 
 def download_yaml(yaml_content, model_name):
@@ -154,26 +163,28 @@ def download_all_data():
     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()
+    average_plot_pil, average_plot_name = plot_average_scores()
+    task_plot_pil, task_plot_name = plot_task_performance()
+    top_models_plot_pil, top_models_plot_name = plot_task_specific_top_models()
+    heatmap_plot_pil, 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)
+        "average_performance": (average_plot_pil, average_plot_name),
+        "task_performance": (task_plot_pil, task_plot_name),
+        "top_models": (top_models_plot_pil, top_models_plot_name),
+        "heatmap": (heatmap_plot_pil, 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 name, (pil_image, filename) in plot_dict.items():
+            image_bytes = io.BytesIO()
+            pil_image.save(image_bytes, format='PNG')
+            image_bytes.seek(0)
+            zf.writestr(filename, image_bytes.read())
+
 
         for model_name in df_full["Model Configuration"].to_list():
             yaml_content = scrape_mergekit_config(model_name)
@@ -184,6 +195,38 @@ def download_all_data():
     
     return zip_buffer, "analysis_data.zip"
 
+def scrape_model_page(model_url):
+    """
+    Scrapes the Hugging Face model page for YAML configuration and other details.
+    """
+    try:
+        # Fetch the model page
+        response = requests.get(model_url)
+        if response.status_code != 200:
+            return f"Error: Unable to fetch the page (Status Code: {response.status_code})"
+        
+        soup = BeautifulSoup(response.text, "html.parser")
+
+        # Extract YAML configuration (usually inside <pre> tags)
+        yaml_config = soup.find("pre")
+        yaml_text = yaml_config.text.strip() if yaml_config else "No YAML configuration found."
+
+        # Extract additional metadata or performance (if available)
+        metadata_section = soup.find("div", class_="metadata")
+        metadata_text = metadata_section.text.strip() if metadata_section else "No metadata found."
+
+        # Return the scraped details
+        return f"**YAML Configuration:**\n{yaml_text}\n\n**Metadata:**\n{metadata_text}"
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+def display_scraped_model_data(model_url):
+    """
+    Displays YAML configuration and metadata for a given model URL.
+    """
+    return scrape_model_page(model_url)
+
 
 # Gradio app
 with gr.Blocks() as demo:
@@ -191,25 +234,25 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         btn1 = gr.Button("Show Average Performance")
-        img1 = gr.Image(type="bytes", label="Average Performance Plot")
+        img1 = gr.Image(type="pil", label="Average Performance Plot",source="upload")
         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="bytes", label="Task Performance Plot")
+        img2 = gr.Image(type="pil", label="Task Performance Plot", source="upload")
         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="bytes", label="Task-Specific Top Models Plot")
+        img3 = gr.Image(type="pil", label="Task-Specific Top Models Plot", source="upload")
         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="bytes", label="Performance Heatmap")
+      heatmap_img = gr.Image(type="pil", label="Performance Heatmap", source="upload")
       heatmap_download = gr.File(label="Download Heatmap")
       btn4.click(plot_heatmap, outputs=[heatmap_img, heatmap_download])
 
@@ -229,6 +272,14 @@ with gr.Blocks() as demo:
         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)
+        
+    # Live scraping feature
+    gr.Markdown("## Live Scraping Features")
+    with gr.Row():
+        url_input = gr.Textbox(label="Enter Hugging Face Model URL", placeholder="https://huggingface.co/<model>")
+        live_scrape_btn = gr.Button("Scrape Model Page")
+        live_scrape_output = gr.Textbox(label="Scraped Data", lines=15)
+        live_scrape_btn.click(display_scraped_model_data, inputs=url_input, outputs=live_scrape_output)
 
 
 demo.launch()