Update app.py

app.py

@@ -4,15 +4,25 @@ import io
 import os
 from pathlib import Path
 from huggingface_hub import HfApi, Repository
+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("")
+    df["label"] = df["label"].fillna("")  # ensure there’s always a label column
     return (
         gr.update(value=df[["text","label"]], visible=True),
         "✅ File uploaded — you can now edit labels."
@@ -29,6 +39,19 @@ 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")
+    plt.tight_layout()
+    return fig
+
 def push_to_hub(repo_name: str, hf_token: str) -> str:
     global df
     try:
@@ -49,13 +72,12 @@ def push_to_hub(repo_name: str, hf_token: str) -> str:
         repo = Repository(
             local_dir=str(local_dir),
             clone_from=repo_name,
-            repo_type="dataset",       # <-- important fix!
+            repo_type="dataset",
             use_auth_token=hf_token
         )
 
         csv_path = local_dir / "data.csv"
         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}"
 
@@ -64,7 +86,7 @@ def push_to_hub(repo_name: str, hf_token: str) -> str:
 
 with gr.Blocks(theme=gr.themes.Default()) as app:
     gr.Markdown("## 🏷️ Label It! Text Annotation Tool")
-    gr.Markdown("Upload a `.csv` with `text` + `label` columns, annotate in-place, then export or publish.")
+    gr.Markdown("Upload a `.csv` with **text** + **label** columns, annotate in-place, then export, visualize, or publish.")
 
     with gr.Row():
         file_input = gr.File(label="📁 Upload CSV", file_types=[".csv"])
@@ -81,17 +103,28 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
     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")
+
+    with gr.Row():
+        model_dropdown = gr.Dropdown(
+            label="🤖 (Future) Auto-Label Model",
+            choices=DEFAULT_MODELS,
+            value=DEFAULT_MODELS[0]
+        )
 
     with gr.Accordion("📦 Push to Hugging Face Hub", open=False):
         repo_input  = gr.Textbox(label="Repo (username/dataset-name)")
-        token_input = gr.Textbox(label="HF Token", type="password")
+        token_input = gr.Textbox(label="🔑 HF Token", type="password")
         push_btn    = gr.Button("🚀 Push")
         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)
+    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)
 
 app.launch()