Update app.py

app.py

@@ -1,13 +1,14 @@
-
+#!/usr/bin/env python3
+"""Enhanced web document annotation tool with modern UI."""
 import hashlib
 import json
 import os
 import uuid
-from collections import defaultdict, deque
+from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from pathlib import Path
-from random import randint, randrange
+from random import sample, shuffle
 
 import gradio as gr
 from datasets import Dataset, load_dataset
@@ -19,9 +20,7 @@ def doc_hash(url: str, text: str) -> str:
 
 
 def filterfunc(x: dict) -> bool:
-    # text length
     if len(x.get("text", "").split()) < 100:
-        # very short content usually means it's not a high quality document
         return False
 
     excluded = {"Promotional/Advertisement", "Machine-Generated", "Images/Videos/Audio", 
@@ -36,15 +35,13 @@ def filterfunc(x: dict) -> bool:
 
 
 class DocLoader:
-    __slots__ = ("queue", "k", "counts", "processed", "total_docs", "_dataset")
+    __slots__ = ("docs", "index", "processed", "_dataset")
 
-    def __init__(self, processed: set[str], k: int = 20):
-        self.queue = deque()
-        self.k = k
-        self.counts = defaultdict(int)
+    def __init__(self, processed: set[str]):
         self.processed = processed
-        self.total_docs = 0
-        self._dataset = None
+        self.index = 0
+        self.docs = []
+        self._dataset = {}
         self._load()
 
     def _load(self):
@@ -54,34 +51,38 @@ class DocLoader:
         ds = ds.filter(filterfunc)
         logger.info(f"Filtered to {len(ds)} documents")
         
-        self._dataset = {}
+        # Build dataset lookup and collect unprocessed docs
+        unprocessed = []
+        
         for idx, doc in enumerate(ds):
             doc_key = doc.get("id", idx)
             doc_with_key = dict(doc)
             doc_with_key["_dataset_key"] = doc_key
             self._dataset[doc_key] = doc_with_key
-        
-        for doc_id, doc in self._dataset.items():
+            
+            # Check if already processed
             url = doc.get("metadata", {}).get("url", doc.get("url", ""))
             h = doc_hash(url, doc.get("text", ""))
             
-            if h in self.processed:
-                continue
-
-            if cat := doc.get("eai_taxonomy", {}).get("document_type_v2", {}).get("primary", {}).get("label"):
-                min_count = min(self.counts.values(), default=0)
-                if self.counts[cat] <= min_count or randrange(self.k) == 0:
-                    self.queue.append(doc)
-                    self.counts[cat] += 1
+            if h not in self.processed:
+                unprocessed.append(doc_with_key)
         
-        self.total_docs = len(self.queue)
-        logger.info(f"Loaded {self.total_docs} documents")
+        logger.info(f"Found {len(unprocessed)} unprocessed documents")
+        
+        # Randomize the order for this session
+        shuffle(unprocessed)
+        self.docs = unprocessed
+        
+        logger.info(f"Loaded {len(self.docs)} documents for this session")
 
     def next(self) -> dict | None:
-        return self.queue.popleft() if self.queue else None
+        if self.index < len(self.docs):
+            doc = self.docs[self.index]
+            self.index += 1
+            return doc
+        return None
 
     def get_by_id(self, doc_id: str | int) -> dict | None:
-        # Handle both string and int lookups
         result = self._dataset.get(doc_id)
         if result is None and isinstance(doc_id, str) and doc_id.isdigit():
             result = self._dataset.get(int(doc_id))
@@ -91,7 +92,7 @@ class DocLoader:
 
     @property
     def remaining(self) -> int:
-        return len(self.queue)
+        return max(0, len(self.docs) - self.index)
 
 
 @dataclass(slots=True)
@@ -99,7 +100,7 @@ class AnnotationStore:
     path: Path
     session_id: str = field(default_factory=lambda: str(uuid.uuid4()))
     buffer: list[dict] = field(default_factory=list)
-    threshold: int = field(default_factory=lambda: randint(20, 30))
+    threshold: int = 25
     processed: set[str] = field(default_factory=set)
     annotations: list[dict] = field(default_factory=list)
     session_stats: dict = field(default_factory=lambda: {
@@ -117,11 +118,7 @@ class AnnotationStore:
                 if rec := self._parse_line(line):
                     self.processed.add(rec["hash"])
                     self.annotations.append(rec)
-                    # Initialize session stats for loaded annotations
-                    if "decision" in rec:
-                        decision = rec["decision"]
-                        if decision not in self.session_stats:
-                            self.session_stats[decision] = 0
+        logger.info(f"Loaded {len(self.processed)} existing annotations")
 
     def _parse_line(self, line: str) -> dict | None:
         try:
@@ -131,6 +128,7 @@ class AnnotationStore:
 
     def add(self, doc_hash: str, decision: str, doc_id: str | int):
         if doc_hash in self.processed:
+            logger.warning(f"Attempted to add already processed document: {doc_hash}")
             return
         
         rec = {
@@ -147,9 +145,10 @@ class AnnotationStore:
         self.annotations.append(rec)
         
         self.session_stats["total"] += 1
-        if decision not in self.session_stats:
-            self.session_stats[decision] = 0
-        self.session_stats[decision] += 1
+        if decision == "selected":
+            self.session_stats["selected"] += 1
+        elif decision == "discarded":
+            self.session_stats["discarded"] += 1
         self.session_stats["decisions"].append((datetime.now(timezone.utc), decision))
         
         if len(self.buffer) >= self.threshold:
@@ -158,7 +157,6 @@ class AnnotationStore:
     def flush(self):
         if not self.buffer or not (token := os.getenv("HF_TOKEN")):
             self.buffer.clear()
-            self.threshold = randint(20, 30)
             return
             
         try:
@@ -170,8 +168,6 @@ class AnnotationStore:
             self.buffer.clear()
         except Exception as e:
             logger.error(f"Push failed: {e}")
-        finally:
-            self.threshold = randint(20, 30)
 
     def get_rate(self) -> float:
         if not self.session_stats["decisions"]:
@@ -191,7 +187,6 @@ store = AnnotationStore(Path("data/annotations.jsonl"))
 loader = DocLoader(store.processed)
 current = loader.next()
 
-
 # Viewer state
 viewer_state = {
     "annotations": [],
@@ -204,9 +199,6 @@ def format_stats() -> str:
     stats = store.session_stats
     rate = store.get_rate()
     
-    selected_count = stats.get('selected', 0)
-    discarded_count = stats.get('discarded', 0)
-    
     return f"""
     <div class="stats-container">
         <div class="stat-item">
@@ -214,17 +206,21 @@ def format_stats() -> str:
             <div class="stat-label">Total Annotated</div>
         </div>
         <div class="stat-item">
-            <div class="stat-value">{selected_count}</div>
+            <div class="stat-value">{stats['selected']}</div>
             <div class="stat-label">Selected</div>
         </div>
         <div class="stat-item">
-            <div class="stat-value">{discarded_count}</div>
+            <div class="stat-value">{stats['discarded']}</div>
             <div class="stat-label">Discarded</div>
         </div>
         <div class="stat-item">
             <div class="stat-value">{rate:.0f}/hr</div>
             <div class="stat-label">Annotation Rate</div>
         </div>
+        <div class="stat-item">
+            <div class="stat-value">{loader.remaining:,}</div>
+            <div class="stat-label">Remaining Docs</div>
+        </div>
     </div>
     """
 
@@ -348,10 +344,8 @@ def update_viewer_filter(filter_value: str):
     viewer_state["filter"] = filter_value
     viewer_state["index"] = 0
     
-    # Get filtered annotations
     viewer_state["annotations"] = store.get_filtered(filter_value)
     
-    # Log for debugging
     logger.info(f"Filter: {filter_value}, Found {len(viewer_state['annotations'])} annotations")
     
     return update_viewer_display()
@@ -380,7 +374,6 @@ def update_viewer_display():
     doc = loader.get_by_id(annotation["id"])
     
     if not doc:
-        # Log the issue for debugging
         logger.warning(f"Document not found for ID: {annotation['id']} (type: {type(annotation['id'])})")
         return (
             "<div class='viewer-error'>Document not found in dataset</div>",
@@ -714,7 +707,6 @@ def build() -> gr.Blocks:
                     viewer_counter = gr.HTML("<div class='viewer-counter'>0 / 0</div>")
                     next_btn = gr.Button("Next →", size="lg")
                 
-                # Initialize viewer
                 filter_dropdown.change(
                     update_viewer_filter,
                     inputs=[filter_dropdown],
@@ -731,7 +723,6 @@ def build() -> gr.Blocks:
                     outputs=[viewer_info, viewer_text, viewer_counter, prev_btn, next_btn]
                 )
                 
-                # Load initial viewer state
                 demo.load(
                     lambda: update_viewer_filter("all"),
                     outputs=[viewer_info, viewer_text, viewer_counter, prev_btn, next_btn]