Upload app.py

app.py

@@ -0,0 +1,382 @@
+import os
+import subprocess
+import signal
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+import gradio as gr
+import tempfile
+import torch
+from datasets import load_dataset
+from tqdm.auto import tqdm
+import re
+import numpy as np
+import gc
+import unicodedata
+from multiprocessing import cpu_count
+from transformers import LlamaTokenizerFast
+import fasttext
+from typing import Tuple, Dict, List
+import json
+import matplotlib.pyplot as plt
+import seaborn as sns
+from datetime import datetime
+import warnings
+from huggingface_hub import HfApi, create_repo, upload_file, snapshot_download, whoami
+from gradio_huggingfacehub_search import HuggingfaceHubSearch
+from pathlib import Path
+from textwrap import dedent
+from scipy import stats
+from apscheduler.schedulers.background import BackgroundScheduler
+
+warnings.filterwarnings('ignore')
+
+# Environment variables
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Global variables for model caching
+MODEL_CACHE_DIR = Path.home() / ".cache" / "ultra_fineweb"
+MODEL_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+MODEL_LOADED = False
+fasttext_model = None
+tokenizer = None
+
+# CSS
+css = """
+.gradio-container {overflow-y: auto;}
+.gr-button-primary {
+    background-color: #ff6b00 !important;
+    border-color: #ff6b00 !important;
+}
+.gr-button-primary:hover {
+    background-color: #ff8534 !important;
+    border-color: #ff8534 !important;
+}
+.gr-button-secondary {
+    background-color: #475467 !important;
+}
+#login-button {
+    background-color: #FFD21E !important;
+    color: #000000 !important;
+}
+"""
+
+# HTML templates
+TITLE = """
+<div style="text-align: center; margin-bottom: 30px;">
+    <h1 style="font-size: 36px; margin-bottom: 10px;">Create your own Dataset Quality Scores, blazingly fast ⚡!</h1>
+    <p style="font-size: 16px; color: #666;">The space takes a HF dataset as input, scores it and provides statistics and quality distribution.</p>
+</div>
+"""
+
+# FIXED: Added `color: #444;` to ensure text is visible on the light background.
+DESCRIPTION = """
+<div style="padding: 20px; background-color: #f0f0f0; border-radius: 10px; margin-bottom: 20px; color: #444;">
+    <h3>📋 How it works:</h3>
+    <ol>
+        <li>Choose a dataset from Hugging Face Hub.</li>
+        <li>The Ultra-FineWeb classifier will score each text sample.</li>
+        <li>View quality distribution and download the scored dataset.</li>
+        <li>Optionally, upload the results to a new repository on your Hugging Face account.</li>
+    </ol>
+    <p><strong>Note:</strong> The first run will download the model (~347MB), which may take a moment.</p>
+</div>
+"""
+
+# --- Helper Functions ---
+def escape(s: str) -> str:
+    """Escape HTML for safe display"""
+    s = str(s).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("\n", "<br/>")
+    return s
+
+def fasttext_preprocess(content: str, tokenizer) -> str:
+    """Preprocess text for FastText model"""
+    if not isinstance(content, str): 
+        return ""
+    content = re.sub(r'\n{3,}', '\n\n', content).lower()
+    content = ''.join(c for c in unicodedata.normalize('NFKD', content) 
+                     if unicodedata.category(c) != 'Mn')
+    token_ids = tokenizer.encode(content, add_special_tokens=False)
+    single_text_list = [tokenizer.decode([token_id]) for token_id in token_ids]
+    content = ' '.join(single_text_list)
+    content = re.sub(r'\n', ' n ', content)
+    content = re.sub(r'\r', '', content)
+    content = re.sub(r'\t', ' ', content)
+    content = re.sub(r' +', ' ', content).strip()
+    return content
+
+def fasttext_infer(norm_content: str, model) -> Tuple[str, float]:
+    """Run FastText inference"""
+    pred_label, pred_prob = model.predict(norm_content)
+    pred_label = pred_label[0]
+    _score = min(pred_prob.tolist()[0], 1.0)
+    if pred_label == "__label__neg":
+        _score = 1 - _score
+    return pred_label, _score
+
+def load_models():
+    """Load models with caching"""
+    global MODEL_LOADED, fasttext_model, tokenizer
+    if MODEL_LOADED:
+        return True
+    
+    try:
+        model_dir = MODEL_CACHE_DIR / "Ultra-FineWeb-classifier"
+        if not model_dir.exists():
+            print("Downloading Ultra-FineWeb-classifier...")
+            snapshot_download(repo_id="openbmb/Ultra-FineWeb-classifier", local_dir=str(model_dir), local_dir_use_symlinks=False)
+        
+        fasttext_path = model_dir / "classifiers" / "ultra_fineweb_en.bin"
+        tokenizer_path = model_dir / "local_tokenizer"
+        
+        if not fasttext_path.exists():
+            raise FileNotFoundError(f"FastText model not found at {fasttext_path}")
+        
+        print("Loading models...")
+        fasttext_model = fasttext.load_model(str(fasttext_path))
+        tokenizer = LlamaTokenizerFast.from_pretrained(str(tokenizer_path) if tokenizer_path.exists() else "meta-llama/Llama-2-7b-hf")
+        
+        MODEL_LOADED = True
+        print("Models loaded successfully!")
+        return True
+    except Exception as e:
+        print(f"Error loading models: {e}")
+        gr.Warning(f"Failed to load models: {e}")
+        return False
+
+def create_quality_plot(scores: List[float], dataset_name: str) -> str:
+    """Create quality distribution plot"""
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
+        output_path = tmpfile.name
+    
+    plt.figure(figsize=(10, 6))
+    sns.histplot(scores, bins=50, kde=True, color='#6B7FD7', edgecolor='black', line_kws={'linewidth': 2, 'color': 'red'})
+    
+    mean_score = np.mean(scores)
+    median_score = np.median(scores)
+    
+    plt.axvline(mean_score, color='green', linestyle='--', linewidth=2, label=f'Mean: {mean_score:.3f}')
+    plt.axvline(median_score, color='orange', linestyle=':', linewidth=2, label=f'Median: {median_score:.3f}')
+    
+    plt.xlabel('Quality Score', fontsize=12)
+    plt.ylabel('Density', fontsize=12)
+    plt.title(f'Quality Score Distribution - {dataset_name}', fontsize=14, fontweight='bold')
+    plt.legend()
+    plt.grid(axis='y', alpha=0.3)
+    plt.xlim(0, 1)
+    
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches='tight')
+    plt.close()
+    
+    return output_path
+
+def process_dataset(
+    model_id: str,
+    dataset_split: str,
+    text_column: str,
+    sample_size: int,
+    batch_size: int,
+    progress=gr.Progress(track_tqdm=True)
+) -> Tuple[str, str, str, str, gr.update, gr.update]:
+    """Process dataset and return results, including visibility updates for UI components."""
+    
+    try:
+        progress(0, desc="Loading models...")
+        if not load_models():
+            raise gr.Error("Failed to load scoring models. Please check the logs.")
+        
+        progress(0.1, desc="Loading dataset...")
+        dataset = load_dataset(model_id, split=dataset_split, streaming=False)
+        
+        if text_column not in dataset.column_names:
+            raise gr.Error(f"Column '{text_column}' not found. Available columns: {', '.join(dataset.column_names)}")
+        
+        total_samples = len(dataset)
+        actual_samples = min(sample_size, total_samples)
+        dataset = dataset.select(range(actual_samples))
+        
+        scores, scored_data = [], []
+        
+        for i in tqdm(range(0, actual_samples, batch_size), desc="Scoring batches"):
+            batch = dataset[i:min(i + batch_size, actual_samples)]
+            for text in batch[text_column]:
+                norm_content = fasttext_preprocess(text, tokenizer)
+                label, score = (0.0, "__label__neg") if not norm_content else fasttext_infer(norm_content, fasttext_model)
+                scores.append(score)
+                scored_data.append({'text': text, 'quality_score': score, 'predicted_label': label})
+        
+        progress(0.9, desc="Generating statistics and plot...")
+        stats_dict = {
+            'dataset_id': model_id,
+            'dataset_split': dataset_split,
+            'processed_samples': actual_samples,
+            'statistics': {
+                'mean': float(np.mean(scores)), 'median': float(np.median(scores)),
+                'std': float(np.std(scores)), 'min': float(np.min(scores)), 'max': float(np.max(scores)),
+                'p90': float(np.percentile(scores, 90)),
+            },
+        }
+        
+        plot_file = create_quality_plot(scores, model_id.split('/')[-1])
+        
+        with tempfile.NamedTemporaryFile(mode='w', suffix=".jsonl", delete=False, encoding='utf-8') as f_out:
+            output_file_path = f_out.name
+            for item in scored_data:
+                f_out.write(json.dumps(item, ensure_ascii=False) + '\n')
+        
+        with tempfile.NamedTemporaryFile(mode='w', suffix=".json", delete=False, encoding='utf-8') as f_stats:
+            stats_file_path = f_stats.name
+            json.dump(stats_dict, f_stats, indent=2)
+            
+        summary_html = f"""
+        <div style="padding: 15px; background-color: #f9f9f9; border-radius: 10px;">
+            <h4>✅ Scoring Completed!</h4>
+            <p><strong>Dataset:</strong> {escape(model_id)}<br>
+            <strong>Processed Samples:</strong> {actual_samples:,}<br>
+            <strong>Mean Score:</strong> {stats_dict['statistics']['mean']:.3f}<br>
+            <strong>Median Score:</strong> {stats_dict['statistics']['median']:.3f}</p>
+        </div>
+        """
+        
+        return summary_html, output_file_path, stats_file_path, plot_file, gr.update(visible=True), gr.update(visible=True)
+            
+    except Exception as e:
+        error_html = f"""
+        <div style="padding: 20px; background-color: #fee; border: 1px solid #d00; border-radius: 10px;">
+            <h4>❌ Error</h4><pre style="white-space: pre-wrap; font-size: 14px;">{escape(e)}</pre>
+        </div>
+        """
+        return error_html, None, None, None, gr.update(visible=False), gr.update(visible=False)
+
+def upload_to_hub(
+    scored_file: str, stats_file: str, plot_file: str, new_dataset_id: str,
+    private: bool, hf_token: str, progress=gr.Progress(track_tqdm=True)
+) -> str:
+    """Upload results to Hugging Face Hub"""
+    if not hf_token: return '❌ <span style="color: red;">Please provide your Hugging Face token.</span>'
+    if not all([scored_file, new_dataset_id]): return '❌ <span style="color: red;">Missing scored file or new dataset ID.</span>'
+    
+    try:
+        progress(0.1, desc="Connecting to Hub...")
+        api = HfApi(token=hf_token)
+        username = whoami(token=hf_token)["name"]
+        repo_id = f"{username}/{new_dataset_id}" if "/" not in new_dataset_id else new_dataset_id
+        
+        progress(0.2, desc=f"Creating repo: {repo_id}")
+        repo_url = create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, private=private, token=hf_token).repo_url
+        
+        progress(0.4, desc="Uploading scored dataset...")
+        upload_file(path_or_fileobj=scored_file, path_in_repo="data/scored_dataset.jsonl", repo_id=repo_id, repo_type="dataset", token=hf_token)
+        
+        progress(0.6, desc="Uploading assets...")
+        if stats_file and os.path.exists(stats_file):
+            upload_file(path_or_fileobj=stats_file, path_in_repo="statistics.json", repo_id=repo_id, repo_type="dataset", token=hf_token)
+        if plot_file and os.path.exists(plot_file):
+            upload_file(path_or_fileobj=plot_file, path_in_repo="quality_distribution.png", repo_id=repo_id, repo_type="dataset", token=hf_token)
+            
+        readme_content = dedent(f"""
+        ---
+        license: apache-2.0
+        ---
+        # Quality-Scored Dataset: {repo_id.split('/')[-1]}
+        This dataset was scored for quality using the [Dataset Quality Scorer Space](https://huggingface.co/spaces/ggml-org/dataset-quality-scorer).
+        ![Quality Distribution](quality_distribution.png)
+        ## Usage
+        ```python
+        from datasets import load_dataset
+        dataset = load_dataset("{repo_id}", split="train")
+        ```
+        """).strip()
+        
+        upload_file(path_or_fileobj=readme_content.encode(), path_in_repo="README.md", repo_id=repo_id, repo_type="dataset", token=hf_token)
+        progress(1.0, "Done!")
+        return f'✅ <span style="color: green;">Successfully uploaded to <a href="{repo_url}" target="_blank">{repo_id}</a></span>'
+        
+    except Exception as e:
+        return f'❌ <span style="color: red;">Upload failed: {escape(e)}</span>'
+
+def create_demo():
+    with gr.Blocks(css=css, title="Dataset Quality Scorer") as demo:
+        gr.HTML(TITLE)
+        gr.HTML(DESCRIPTION)
+        
+        gr.Markdown("### 1. Configure & Score Dataset")
+        with gr.Row():
+            with gr.Column(scale=3):
+                dataset_search = HuggingfaceHubSearch(label="Hub Dataset ID", search_type="dataset", value="roneneldan/TinyStories")
+                text_column = gr.Textbox(label="Text Column Name", value="text", info="The column containing the text to score.")
+            with gr.Column(scale=2):
+                dataset_split = gr.Dropdown(["train", "validation", "test"], label="Split", value="train")
+                with gr.Row():
+                    sample_size = gr.Number(label="Sample Size", value=1000, minimum=100, step=100, info="Max samples.")
+                    batch_size = gr.Number(label="Batch Size", value=32, minimum=1, step=1, info="Processing batch.")
+
+        with gr.Row():
+            clear_btn = gr.Button("Clear", variant="secondary")
+            process_btn = gr.Button("🚀 Start Scoring", variant="primary", size="lg")
+        
+        # --- Results and Upload Sections (Initially Hidden) ---
+        with gr.Accordion("✅ Results", open=True, visible=False) as results_accordion:
+            gr.Markdown("### 2. Review Results")
+            with gr.Row():
+                with gr.Column(scale=2):
+                    summary_output = gr.HTML(label="Summary")
+                with gr.Column(scale=1):
+                    plot_output = gr.Image(label="Quality Distribution", show_label=True)
+            with gr.Row():
+                scored_file_output = gr.File(label="📄 Download Scored Dataset (.jsonl)", type="filepath")
+                stats_file_output = gr.File(label="📊 Download Statistics (.json)", type="filepath")
+
+        with gr.Accordion("☁️ Upload to Hub", open=False, visible=False) as upload_accordion:
+            gr.Markdown("### 3. (Optional) Upload to Hugging Face Hub")
+            hf_token_input = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_...", value=HF_TOKEN or "", info="Your HF token with 'write' permissions.")
+            new_dataset_id = gr.Textbox(label="New Dataset Name", placeholder="my-scored-dataset", info="Will be created under your username.")
+            private_checkbox = gr.Checkbox(label="Make dataset private", value=False)
+            upload_btn = gr.Button("📤 Upload to Hub", variant="primary")
+            upload_status = gr.HTML()
+
+        # --- Event Handlers ---
+        def clear_form():
+            return "roneneldan/TinyStories", "train", "text", 1000, 32, None, None, None, None, gr.update(visible=False), gr.update(visible=False), ""
+        
+        clear_btn.click(
+            fn=clear_form,
+            outputs=[
+                dataset_search, dataset_split, text_column, sample_size, batch_size,
+                summary_output, scored_file_output, stats_file_output, plot_output,
+                results_accordion, upload_accordion, upload_status
+            ]
+        )
+        
+        process_btn.click(
+            fn=process_dataset,
+            inputs=[dataset_search, dataset_split, text_column, sample_size, batch_size],
+            outputs=[summary_output, scored_file_output, stats_file_output, plot_output, results_accordion, upload_accordion]
+        )
+        
+        upload_btn.click(
+            fn=upload_to_hub,
+            inputs=[scored_file_output, stats_file_output, plot_output, new_dataset_id, private_checkbox, hf_token_input],
+            outputs=[upload_status]
+        )
+    return demo
+
+# --- App Execution ---
+demo = create_demo()
+
+if os.environ.get("SPACE_ID"):
+    def restart_space():
+        if HF_TOKEN:
+            try:
+                print("Scheduler: Triggering space restart...")
+                api = HfApi()
+                api.restart_space(repo_id=os.environ["SPACE_ID"], token=HF_TOKEN)
+            except Exception as e:
+                print(f"Scheduler: Failed to restart space: {e}")
+                
+    scheduler = BackgroundScheduler()
+    scheduler.add_job(restart_space, "interval", hours=6)
+    scheduler.start()
+    print("Background scheduler for periodic restarts is active.")
+
+if __name__ == "__main__":
+    demo.queue().launch(debug=False, show_api=False)