Update app.py

app.py

@@ -9,20 +9,12 @@ import matplotlib.pyplot as plt
 # Global DataFrame
 df = pd.DataFrame()
 
-# List of free, recommended models (for future auto-labeling)
-DEFAULT_MODELS = [
-    "mistralai/Mistral-7B-Instruct-v0.2",
-    "HuggingFaceH4/zephyr-7b-beta",
-    "tiiuae/falcon-rw-1b",
-    "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-]
-
 def upload_csv(file):
     global df
     df = pd.read_csv(file.name)
     if "text" not in df.columns or "label" not in df.columns:
         return gr.update(visible=False), "❌ CSV must contain ‘text’ and ‘label’ columns."
-    df["label"] = df["label"].fillna("")  # ensure there’s always a label column
+    df["label"] = df["label"].fillna("")
     return (
         gr.update(value=df[["text","label"]], visible=True),
         "✅ File uploaded — you can now edit labels."
@@ -39,34 +31,50 @@ def download_csv():
     df.to_csv(out_path, index=False)
     return out_path
 
-def visualize_distribution():
-    global df
-    if df.empty or "label" not in df.columns:
-        return None
-    counts = df["label"].value_counts()
-    fig, ax = plt.subplots()
-    counts.plot(kind="bar", ax=ax)
-    ax.set_title("Label Distribution")
-    ax.set_xlabel("Label")
-    ax.set_ylabel("Count")
+def create_distribution_figure(df_input):
+    counts = df_input["label"].value_counts().sort_values(ascending=False)
+    labels = counts.index.tolist()
+    values = counts.values.tolist()
+    
+    fig, (ax_table, ax_bar) = plt.subplots(
+        nrows=1, ncols=2,
+        gridspec_kw={"width_ratios": [1, 2]},
+        figsize=(8, max(2, len(labels) * 0.3))
+    )
+    
+    # Table
+    ax_table.axis("off")
+    table_data = [[lab, cnt] for lab, cnt in zip(labels, values)]
+    tbl = ax_table.table(cellText=table_data, colLabels=["Label","Count"], loc="center")
+    tbl.auto_set_font_size(False)
+    tbl.set_fontsize(10)
+    tbl.scale(1, 1.5)
+    
+    # Bar chart
+    ax_bar.barh(labels, values)
+    ax_bar.invert_yaxis()
+    ax_bar.set_xlabel("Count")
+    ax_bar.set_ylabel("")
+    
     plt.tight_layout()
     return fig
 
+def visualize_and_download_chart():
+    global df
+    fig = create_distribution_figure(df)
+    chart_path = "label_distribution.png"
+    fig.savefig(chart_path, dpi=150)
+    return fig, chart_path
+
 def push_to_hub(repo_name: str, hf_token: str) -> str:
     global df
     try:
         api = HfApi()
-        api.create_repo(
-            repo_id=repo_name,
-            token=hf_token,
-            repo_type="dataset",
-            exist_ok=True
-        )
+        api.create_repo(repo_id=repo_name, token=hf_token, repo_type="dataset", exist_ok=True)
 
         local_dir = Path(f"./{repo_name.replace('/', '_')}")
         if local_dir.exists():
-            for child in local_dir.iterdir():
-                child.unlink()
+            for child in local_dir.iterdir(): child.unlink()
             local_dir.rmdir()
 
         repo = Repository(
@@ -80,7 +88,6 @@ def push_to_hub(repo_name: str, hf_token: str) -> str:
         df.to_csv(csv_path, index=False)
         repo.push_to_hub(commit_message="📑 Update annotated data")
         return f"🚀 Pushed to https://huggingface.co/datasets/{repo_name}"
-
     except Exception as e:
         return f"❌ Push failed: {e}"
 
@@ -92,27 +99,19 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
         file_input = gr.File(label="📁 Upload CSV", file_types=[".csv"])
         upload_btn = gr.Button("Upload")
 
-    df_table = gr.Dataframe(
-        headers=["text","label"],
-        label="📝 Editable Table",
-        interactive=True,
-        visible=False
-    )
+    df_table = gr.Dataframe(headers=["text","label"], label="📝 Editable Table",
+                            interactive=True, visible=False)
     status = gr.Textbox(label="Status", interactive=False)
 
     with gr.Row():
-        save_btn     = gr.Button("💾 Save")
-        download_btn = gr.Button("⬇️ Download CSV")
-        visualize_btn= gr.Button("📊 Visualize Distribution")
-        download_out = gr.File(label="📥 Downloaded File")
-        viz_out      = gr.Plot(label="Label Distribution")
+        save_btn        = gr.Button("💾 Save")
+        download_btn    = gr.Button("⬇️ Download CSV")
+        download_out    = gr.File(label="📥 Downloaded File")
 
     with gr.Row():
-        model_dropdown = gr.Dropdown(
-            label="🤖 (Future) Auto-Label Model",
-            choices=DEFAULT_MODELS,
-            value=DEFAULT_MODELS[0]
-        )
+        visualize_btn   = gr.Button("📊 Visualize Distribution")
+        chart_plot      = gr.Plot(label="Label Distribution")
+        download_chart  = gr.File(label="📥 Download Chart")
 
     with gr.Accordion("📦 Push to Hugging Face Hub", open=False):
         repo_input  = gr.Textbox(label="Repo (username/dataset-name)")
@@ -121,10 +120,10 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
         push_status = gr.Textbox(label="Push Status", interactive=False)
 
     # Bind events
-    upload_btn.click(upload_csv,     inputs=file_input,              outputs=[df_table, status])
-    save_btn.click(  save_changes,   inputs=df_table,                outputs=status)
+    upload_btn.click(upload_csv, inputs=file_input, outputs=[df_table, status])
+    save_btn.click(save_changes, inputs=df_table, outputs=status)
     download_btn.click(download_csv, outputs=download_out)
-    visualize_btn.click(visualize_distribution, outputs=viz_out)
-    push_btn.click(  push_to_hub,    inputs=[repo_input, token_input], outputs=push_status)
+    visualize_btn.click(visualize_and_download_chart, outputs=[chart_plot, download_chart])
+    push_btn.click(push_to_hub, inputs=[repo_input, token_input], outputs=push_status)
 
 app.launch()