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 = """

Create your own Dataset Quality Scores, blazingly fast ⚡!

The space takes a HF dataset as input, scores it and provides statistics and quality distribution.

""" # FIXED: Added `color: #444;` to ensure text is visible on the light background. DESCRIPTION = """

📋 How it works:

  1. Choose a dataset from Hugging Face Hub.
  2. The Ultra-FineWeb classifier will score each text sample.
  3. View quality distribution and download the scored dataset.
  4. Optionally, upload the results to a new repository on your Hugging Face account.

Note: The first run will download the model (~347MB), which may take a moment.

""" # --- Helper Functions --- def escape(s: str) -> str: """Escape HTML for safe display""" s = str(s).replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """).replace("\n", "
") 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"""

✅ Scoring Completed!

Dataset: {escape(model_id)}
Processed Samples: {actual_samples:,}
Mean Score: {stats_dict['statistics']['mean']:.3f}
Median Score: {stats_dict['statistics']['median']:.3f}

""" 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"""

❌ Error

{escape(e)}
""" 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 '❌ Please provide your Hugging Face token.' if not all([scored_file, new_dataset_id]): return '❌ Missing scored file or new dataset ID.' 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'✅ Successfully uploaded to {repo_id}' except Exception as e: return f'❌ Upload failed: {escape(e)}' 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)