Upload app.py with huggingface_hub

app.py

@@ -1,7 +1,7 @@
 """
 Hunter Omega — Live NIAH Demo
 One-click needle in a haystack test from 8K to 12M tokens.
-Hosted on HuggingFace Spaces.
+Shows the needle BEFORE querying so users can verify it's not rigged.
 """
 import gradio as gr
 import requests
@@ -13,7 +13,6 @@ import os
 API_URL = os.environ.get("HUNTER_API")
 MODEL = "hunter-omega-12m-fast"
 
-# Diverse filler paragraphs (no repetition)
 FILLERS = [
     "The evolution of distributed systems has fundamentally changed how we architect software. Microservices emerged as a response to monolithic limitations, enabling teams to deploy independently and scale specific components. Service meshes handle the complexity of inter-service communication, providing observability and traffic management.",
     "Machine learning pipelines require careful orchestration from data ingestion through model deployment. Feature stores centralize feature computation and serving, ensuring consistency between training and inference. MLOps practices bring version control and automated testing to model development workflows.",
@@ -32,9 +31,14 @@ FILLERS = [
     "Robotics integrates sensing planning and actuation for autonomous interaction. Simultaneous localization and mapping builds environment models while tracking position. Motion planning finds collision-free paths through configuration spaces.",
 ]
 
+SIZE_MAP = {
+    "8K": 8000, "32K": 32000, "64K": 64000, "128K": 128000,
+    "256K": 256000, "512K": 512000, "1M": 1000000,
+    "1.5M": 1500000, "3M": 3000000, "6M": 6000000, "12M": 12000000,
+}
+
 
 def generate_needle():
-    """Generate a random unique passphrase that can't be guessed."""
     code = ''.join(random.choices(string.ascii_uppercase, k=4))
     nums = ''.join(random.choices(string.digits, k=4))
     word = random.choice(["ALPHA", "BRAVO", "DELTA", "ECHO", "FOXTROT", "GOLF", "HOTEL",
@@ -44,7 +48,6 @@ def generate_needle():
 
 
 def generate_haystack(target_tokens):
-    """Generate diverse filler text of target token count."""
     words = []
     target_words = target_tokens * 3 // 4
     i = 0
@@ -54,178 +57,138 @@ def generate_haystack(target_tokens):
     return words[:target_words]
 
 
-def run_niah_native(context_tokens):
-    """Run NIAH test within native 1M window (direct prompt)."""
+def run_test(context_size):
+    """Streaming test: shows needle first, then queries, then shows result."""
+    tokens = SIZE_MAP.get(context_size, 128000)
     needle_code = generate_needle()
     needle = f"The classified authorization code for Operation Starfall is {needle_code}."
     query = "What is the classified authorization code for Operation Starfall? State ONLY the code."
+    depth = random.randint(10, 90) if tokens <= 1000000 else random.randint(55, 95)
 
     # Generate haystack
-    words = generate_haystack(context_tokens)
-    depth = random.randint(10, 90)
+    words = generate_haystack(tokens)
     insert_pos = int(len(words) * depth / 100)
     words.insert(insert_pos, needle)
-    haystack = " ".join(words)
-
-    actual_words = len(haystack.split())
-    actual_tokens_est = actual_words * 4 // 3
-
-    # Query the model
-    t0 = time.time()
-    try:
-        r = requests.post(API_URL, json={
-            "model": MODEL,
-            "messages": [{"role": "user", "content": haystack + "\n\n" + query}],
-            "max_tokens": 50,
-            "temperature": 0,
-        }, timeout=300)
-        elapsed = time.time() - t0
-
-        if r.status_code != 200:
-            return f"Server returned {r.status_code}. Server may be busy, please retry.", "", "", ""
-
-        d = r.json()
-        answer = d["choices"][0]["message"]["content"]
-        found = needle_code.lower() in answer.lower()
+    actual_words = len(words)
+    actual_tokens = actual_words * 4 // 3
 
-        status = "FOUND" if found else "NOT FOUND"
-        result_text = f"""## Result: {"✅ NEEDLE FOUND" if found else "❌ NEEDLE NOT FOUND"}
+    # STEP 1: Show what we're about to test
+    setup_text = f"""## Step 1: Test Setup
 
-**Context:** {actual_tokens_est:,} tokens (~{actual_words:,} words)
-**Needle depth:** {depth}%
-**Query time:** {elapsed:.1f}s
-**Model answer:** {answer[:200]}"""
+**Context size:** {actual_tokens:,} tokens (~{actual_words:,} words)
+**Generated needle:** `{needle_code}`
+**Full needle sentence:** "{needle}"
+**Hidden at:** {depth}% depth (word position {insert_pos:,} of {actual_words:,})
+**Query:** "{query}"
 
-        needle_info = f"""**Hidden needle:** "{needle}"
-**Inserted at:** {depth}% depth (position {insert_pos:,} of {actual_words:,} words)
-**Passphrase:** {needle_code}"""
-
-        return result_text, needle_info, status, f"{elapsed:.1f}s"
-
-    except requests.exceptions.Timeout:
-        return "Request timed out. Server may be processing other requests. Please retry.", "", "TIMEOUT", ""
-    except Exception as e:
-        return f"Connection error. Server may be busy. Please retry in a moment.", "", "ERROR", ""
-
-
-def run_niah_extended(context_tokens):
-    """Run NIAH test beyond 1M using the overflow pipeline."""
-    needle_code = generate_needle()
-    needle = f"The classified authorization code for Operation Starfall is {needle_code}."
-    query = "What is the classified authorization code for Operation Starfall? State ONLY the code."
-
-    # Generate document
-    words = generate_haystack(context_tokens)
-    depth = random.randint(55, 95)  # Always in overflow portion
-    insert_pos = int(len(words) * depth / 100)
-    words.insert(insert_pos, needle)
-    doc = " ".join(words)
-
-    actual_words = len(doc.split())
-    actual_tokens_est = actual_words * 4 // 3
-
-    user_id = f"demo_{random.randint(100000, 999999)}"
-    api_base = API_URL.replace("/v1/chat/completions", "")
-
-    # Upload document
-    t0 = time.time()
-    try:
-        r = requests.post(f"{api_base}/v1/documents", json={
-            "user_id": user_id,
-            "text": doc,
-        }, timeout=1800)
-        upload_time = time.time() - t0
-
-        if r.status_code != 200:
-            return f"Upload failed ({r.status_code}). Server may be busy.", "", "", ""
-
-        info = r.json()
-        chunks = info.get("overflow_chunks", 0)
-
-        # Query
-        t1 = time.time()
-        r = requests.post(API_URL, json={
-            "model": MODEL,
-            "messages": [{"role": "user", "content": query}],
-            "max_tokens": 50,
-            "temperature": 0,
-            "user": user_id,
-        }, timeout=600)
-        query_time = time.time() - t1
-        total_time = time.time() - t0
-
-        # Cleanup
-        requests.delete(f"{api_base}/v1/documents/{user_id}", timeout=10)
-
-        if r.status_code != 200:
-            return f"Query failed ({r.status_code}). Server may be busy.", "", "", ""
-
-        d = r.json()
-        answer = d["choices"][0]["message"]["content"]
-        found = needle_code.lower() in answer.lower()
-
-        result_text = f"""## Result: {"✅ NEEDLE FOUND" if found else "❌ NEEDLE NOT FOUND"}
-
-**Context:** {actual_tokens_est:,} tokens (~{actual_words:,} words)
-**Overflow chunks:** {chunks}
-**Needle depth:** {depth}%
-**Index time:** {upload_time:.1f}s
-**Query time:** {query_time:.1f}s
-**Total time:** {total_time:.1f}s
-**Model answer:** {answer[:200]}"""
-
-        needle_info = f"""**Hidden needle:** "{needle}"
-**Inserted at:** {depth}% depth (position {insert_pos:,} of {actual_words:,} words)
-**Passphrase:** {needle_code}"""
-
-        return result_text, needle_info, "FOUND" if found else "NOT FOUND", f"{total_time:.1f}s"
-
-    except requests.exceptions.Timeout:
-        requests.delete(f"{api_base}/v1/documents/{user_id}", timeout=10)
-        return "Request timed out. Large documents take time to index. Please retry.", "", "TIMEOUT", ""
-    except Exception as e:
-        return f"Connection error. Server may be busy.", "", "ERROR", ""
+The needle is now embedded in {actual_tokens:,} tokens of diverse technical text.
+The model has never seen this passphrase before. It was just generated randomly.
 
+---
 
-def run_test(context_size):
-    """Main test function called by Gradio."""
-    size_map = {
-        "8K": 8000,
-        "32K": 32000,
-        "64K": 64000,
-        "128K": 128000,
-        "256K": 256000,
-        "512K": 512000,
-        "1M": 1000000,
-        "1.5M": 1500000,
-        "3M": 3000000,
-        "6M": 6000000,
-        "12M": 12000000,
-    }
-
-    tokens = size_map.get(context_size, 128000)
+## Step 2: Querying model..."""
 
     if tokens <= 1000000:
-        return run_niah_native(tokens)
+        # Native window: direct prompt
+        haystack = " ".join(words)
+        yield setup_text, "", "QUERYING...", ""
+
+        t0 = time.time()
+        try:
+            r = requests.post(API_URL, json={
+                "model": MODEL,
+                "messages": [{"role": "user", "content": haystack + "\n\n" + query}],
+                "max_tokens": 50,
+                "temperature": 0,
+            }, timeout=300)
+            elapsed = time.time() - t0
+
+            if r.status_code != 200:
+                yield setup_text, "", "ERROR", f"Server returned {r.status_code}"
+                return
+
+            answer = r.json()["choices"][0]["message"]["content"]
+            found = needle_code.lower() in answer.lower()
+
+        except requests.exceptions.Timeout:
+            yield setup_text, "", "TIMEOUT", "Server busy, please retry"
+            return
+        except Exception:
+            yield setup_text, "", "ERROR", "Connection error, please retry"
+            return
     else:
-        return run_niah_extended(tokens)
-
-
-# Build Gradio interface
-with gr.Blocks(title="Hunter Omega — Live NIAH Demo", theme=gr.themes.Soft()) as demo:
+        # Overflow pipeline
+        doc = " ".join(words)
+        user_id = f"demo_{random.randint(100000, 999999)}"
+        api_base = API_URL.replace("/v1/chat/completions", "")
+
+        yield setup_text + "\n\n*Indexing large document...*", "", "INDEXING...", ""
+
+        t0 = time.time()
+        try:
+            r = requests.post(f"{api_base}/v1/documents", json={
+                "user_id": user_id, "text": doc,
+            }, timeout=1800)
+            index_time = time.time() - t0
+            info = r.json()
+            chunks = info.get("overflow_chunks", 0)
+
+            yield setup_text + f"\n\n*Indexed {chunks} chunks in {index_time:.0f}s. Querying...*", "", "QUERYING...", ""
+
+            t1 = time.time()
+            r = requests.post(API_URL, json={
+                "model": MODEL,
+                "messages": [{"role": "user", "content": query}],
+                "max_tokens": 50,
+                "temperature": 0,
+                "user": user_id,
+            }, timeout=600)
+            elapsed = time.time() - t0
+            answer = r.json()["choices"][0]["message"]["content"]
+            found = needle_code.lower() in answer.lower()
+
+            requests.delete(f"{api_base}/v1/documents/{user_id}", timeout=10)
+
+        except requests.exceptions.Timeout:
+            requests.delete(f"{api_base}/v1/documents/{user_id}", timeout=10)
+            yield setup_text, "", "TIMEOUT", "Large doc timed out, please retry"
+            return
+        except Exception:
+            yield setup_text, "", "ERROR", "Connection error"
+            return
+
+    # STEP 3: Show result
+    result_text = f"""## Step 3: Result {"✅ NEEDLE FOUND" if found else "❌ NOT FOUND"}
+
+**Expected:** `{needle_code}`
+**Model returned:** `{answer.strip()}`
+**Match:** {"YES" if found else "NO"}
+**Time:** {elapsed:.1f}s
+
+{"The model correctly retrieved the randomly generated passphrase from " + f"{actual_tokens:,} tokens of text." if found else "The model did not find the needle."}"""
+
+    full_output = setup_text + "\n\n---\n\n" + result_text
+
+    yield full_output, "", "FOUND" if found else "MISSED", f"{elapsed:.1f}s"
+
+
+# Build interface
+with gr.Blocks(title="Hunter Omega — Live NIAH Demo") as demo:
     gr.Markdown("""
-# Hunter Omega — Live Needle in a Haystack Demo
+# Hunter Omega — Live Needle in a Haystack
 
-**Test it yourself.** Pick a context length, click the button. The app generates a random document,
-hides a random passphrase at a random depth, and the model finds it.
+**How this works:**
+1. You pick a context length
+2. The app generates a random passphrase and hides it at a random depth
+3. You see the passphrase and its position BEFORE the model runs
+4. The model searches the document and returns what it found
+5. You verify it matches
 
-You control nothing about the needle. It's random every time. Run it as many times as you want.
+Nothing is pre-arranged. The passphrase is generated fresh every time. Run it as many times as you want.
 
-**8K to 1M:** Native full attention window. Instant results.
-**1.5M to 12M:** Extended context via overflow pipeline. Larger sizes take longer to index.
+**8K to 1M:** Native full attention. **1.5M to 12M:** Extended context engine.
 
-*Running on a single GPU cluster. If the server is busy, responses may be delayed. This is a capacity
-limitation, not an accuracy issue.*
+*Running on a single GPU cluster. If busy, responses may queue. This is capacity, not accuracy.*
     """)
 
     with gr.Row():
@@ -236,33 +199,28 @@ limitation, not an accuracy issue.*
         )
         run_btn = gr.Button("Find the Needle", variant="primary", size="lg")
 
-    with gr.Row():
-        with gr.Column():
-            result_output = gr.Markdown(label="Result")
-        with gr.Column():
-            needle_output = gr.Markdown(label="Hidden Needle (revealed after test)")
+    output = gr.Markdown(label="Test Progress")
 
     with gr.Row():
-        status_output = gr.Textbox(label="Status", interactive=False)
-        time_output = gr.Textbox(label="Time", interactive=False)
+        status_box = gr.Textbox(label="Status", interactive=False)
+        time_box = gr.Textbox(label="Time", interactive=False)
+
+    hidden = gr.Textbox(visible=False)
 
     gr.Markdown("""
 ---
-**Estimated wait times:**
-
-| Size | Index Time | Query Time | Total |
-|:----:|:----------:|:----------:|:-----:|
-| 8K - 1M | None (native) | 1-90s | 1-90s |
-| 1.5M | ~1 min | ~16s | ~1.5 min |
-| 3M | ~3 min | ~16s | ~3.5 min |
-| 6M | ~5 min | ~16s | ~6 min |
-| 12M | ~13 min | ~90s | ~14 min |
+**Estimated times:**
 
-If multiple users are testing, requests queue. The model processes them in order. Your request will complete, it may just take longer.
+| Size | Time |
+|:----:|:----:|
+| 8K - 128K | 2-8s |
+| 256K - 512K | 10-30s |
+| 1M | 60-90s |
+| 1.5M - 3M | 1-4 min |
+| 6M | 5-6 min |
+| 12M | 13-15 min |
 
-**How this works:** A random passphrase is generated and hidden at a random depth in diverse technical
-filler text. The model is asked to retrieve the exact passphrase. No keyword hints in the query.
-The passphrase and its location are revealed after the test so you can verify.
+If multiple users are testing, requests queue behind each other.
 
 [Full benchmarks and methodology](https://github.com/SovNodeAI/hunter-omega-benchmarks)
     """)
@@ -270,7 +228,7 @@ The passphrase and its location are revealed after the test so you can verify.
     run_btn.click(
         fn=run_test,
         inputs=[context_dropdown],
-        outputs=[result_output, needle_output, status_output, time_output],
+        outputs=[output, hidden, status_box, time_box],
     )