Update app.py

app.py

@@ -14,7 +14,7 @@ import unicodedata
 from multiprocessing import cpu_count
 from transformers import LlamaTokenizerFast
 import fasttext
-from typing import Tuple, Dict, List
+from typing import Tuple, Dict, List, Generator
 import json
 import matplotlib.pyplot as plt
 import seaborn as sns
@@ -50,13 +50,6 @@ css = """
     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
@@ -67,44 +60,32 @@ TITLE = """
 </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>
+# Switched to Markdown for better theme compatibility (dark/light mode)
+DESCRIPTION_MD = """
+### 📋 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("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("\n", "<br/>")
-    return s
+    return str(s).replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """).replace("\n", "<br/>")
 
 def fasttext_preprocess(content: str, tokenizer) -> str:
-    """Preprocess text for FastText model"""
-    if not isinstance(content, str): 
-        return ""
+    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')
+    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
+    content = ' '.join([tokenizer.decode([token_id]) for token_id in token_ids])
+    content = re.sub(r'\n', ' n ', content).replace('\r', '').replace('\t', ' ')
+    return re.sub(r' +', ' ', content).strip()
 
 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)
@@ -113,62 +94,38 @@ def fasttext_infer(norm_content: str, model) -> Tuple[str, float]:
     return pred_label, _score
 
 def load_models():
-    """Load models with caching"""
     global MODEL_LOADED, fasttext_model, tokenizer
-    if MODEL_LOADED:
-        return True
-    
+    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)
-    
+    sns.histplot(scores, bins=50, kde=True, color='#6B7FD7', edgecolor='black')
+    mean_score, median_score = np.mean(scores), 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.xlabel('Quality Score'); plt.ylabel('Density')
+    plt.title(f'Quality Score Distribution - {dataset_name}', fontweight='bold')
+    plt.legend(); plt.grid(axis='y', alpha=0.3); plt.xlim(0, 1)
+    plt.tight_layout(); plt.savefig(output_path, dpi=150)
     plt.close()
-    
     return output_path
 
+# UPDATED: This function is now a generator to yield live log updates.
 def process_dataset(
     model_id: str,
     dataset_split: str,
@@ -176,82 +133,80 @@ def process_dataset(
     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."""
-    
+) -> Generator:
+    log_text = ""
+    # Helper to update and yield log messages
+    def update_log(msg):
+        nonlocal log_text
+        timestamp = datetime.now().strftime('%H:%M:%S')
+        log_text += f"[{timestamp}] {msg}\n"
+        # Yield updates for the log, keep other components hidden/empty
+        return (log_text, None, None, None, None, gr.update(visible=False), gr.update(visible=False))
+
     try:
-        progress(0, desc="Loading models...")
+        yield update_log("Starting process...")
+        
+        yield update_log("Loading scoring models...")
         if not load_models():
-            raise gr.Error("Failed to load scoring models. Please check the logs.")
+            raise gr.Error("Failed to load scoring models. Please check logs.")
+        yield update_log("Models loaded successfully.")
         
-        progress(0.1, desc="Loading dataset...")
+        yield update_log(f"Loading dataset '{model_id}' split '{dataset_split}'...")
         dataset = load_dataset(model_id, split=dataset_split, streaming=False)
+        yield update_log("Dataset loaded.")
         
         if text_column not in dataset.column_names:
-            raise gr.Error(f"Column '{text_column}' not found. Available columns: {', '.join(dataset.column_names)}")
+            raise gr.Error(f"Column '{text_column}' not found. Available: {', '.join(dataset.column_names)}")
         
-        total_samples = len(dataset)
-        actual_samples = min(sample_size, total_samples)
+        actual_samples = min(sample_size, len(dataset))
         dataset = dataset.select(range(actual_samples))
         
+        yield update_log(f"Starting to score {actual_samples:,} 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)
+                label, score = fasttext_infer(norm_content, fasttext_model) if norm_content else ("__label__neg", 0.0)
                 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)),
-            },
-        }
+        yield update_log("Scoring complete. Generating results and plot...")
+        stats_dict = {'dataset_id': model_id, 'processed_samples': actual_samples, 'statistics': {'mean': float(np.mean(scores)), 'median': float(np.median(scores))}}
         
         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('w', suffix=".jsonl", delete=False, encoding='utf-8') as f:
+            output_file_path = f.name
+            for item in scored_data: f.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)
+        with tempfile.NamedTemporaryFile('w', suffix=".json", delete=False, encoding='utf-8') as f:
+            stats_file_path = f.name
+            json.dump(stats_dict, f, 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>
+        summary_md = f"""
+        #### ✅ Scoring Completed!
+        - **Dataset:** `{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)
-            
+        yield update_log("Process finished successfully!")
+        
+        # Final return with all components visible and populated
+        yield (log_text, summary_md, 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)
+        error_log = update_log(f"ERROR: {e}")[0]
+        error_summary_md = f"### ❌ Error\n```\n{escape(str(e))}\n```"
+        yield (error_log, error_summary_md, None, None, None, gr.update(visible=True), gr.update(visible=False))
 
+# This function remains the same
 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>'
     
@@ -264,95 +219,99 @@ def upload_to_hub(
         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...")
+        progress(0.4, desc="Uploading files...")
         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()
+        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>'
+        return f'❌ <span style="color: red;">Upload failed: {escape(str(e))}</span>'
+
 
 def create_demo():
     with gr.Blocks(css=css, title="Dataset Quality Scorer") as demo:
         gr.HTML(TITLE)
-        gr.HTML(DESCRIPTION)
+        gr.Markdown(DESCRIPTION_MD)
         
-        gr.Markdown("### 1. Configure & Score Dataset")
         with gr.Row():
             with gr.Column(scale=3):
+                gr.Markdown("### 1. Configure Dataset")
                 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.")
+                text_column = gr.Textbox(label="Text Column Name", value="text")
             with gr.Column(scale=2):
+                gr.Markdown("### 2. Configure Scoring")
                 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.")
+                    sample_size = gr.Number(label="Sample Size", value=1000, minimum=100, step=100)
+                    batch_size = gr.Number(label="Batch Size", value=32, minimum=1, step=1)
 
+        live_log = gr.Textbox(label="Live Log", interactive=False, lines=8, max_lines=20)
+        
         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.Group(visible=False) as results_group:
+            gr.Markdown("--- \n ### 3. Review Results")
             with gr.Row():
-                with gr.Column(scale=2):
-                    summary_output = gr.HTML(label="Summary")
+                with gr.Column(scale=1):
+                    summary_output = gr.Markdown(label="Summary")
+                    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.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.")
+        with gr.Group(visible=False) as upload_group:
+            gr.Markdown("--- \n ### 4. (Optional) Upload to Hugging Face Hub")
+            hf_token_input = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_...", value=HF_TOKEN or "")
+            new_dataset_id = gr.Textbox(label="New Dataset Name", placeholder="my-scored-dataset")
             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), ""
+            return "roneneldan/TinyStories", "train", "text", 1000, 32, "", None, None, None, None, gr.update(visible=False), gr.update(visible=False), ""
+        
+        outputs_list = [
+            live_log, summary_output, scored_file_output, stats_file_output, plot_output,
+            results_group, upload_group
+        ]
+        
+        process_btn.click(
+            fn=process_dataset,
+            inputs=[dataset_search, dataset_split, text_column, sample_size, batch_size],
+            outputs=outputs_list
+        )
         
         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
+                live_log, summary_output, scored_file_output, stats_file_output, plot_output,
+                results_group, upload_group, 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],
@@ -363,20 +322,5 @@ def create_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)