Upload experiments/needle_in_haystack.py

experiments/needle_in_haystack.py

@@ -1,8 +1,9 @@
 """
 Experiment 2: Needle in Haystack (text)
-Tests retrieval of a secret code hidden at varying depths in filler text.
-FIXED: Increased default context to 1500 sentences to stress 1.5B model attention.
-Also adds decoy codes to prevent trivial keyword-only retrieval.
+Tests retrieval of a fact hidden at varying depths in filler text.
+FIXED: 2000-sentence haystacks + entity-overlap distractors to prevent keyword-only
+retrieval. The target entities (person, item, location) each appear in multiple
+sentences, forcing the model to attend to the right combination.
 """
 import logging
 import os
@@ -19,66 +20,81 @@ from src.utils import ensure_dir, save_jsonl, save_json
 
 logger = logging.getLogger(__name__)
 
+# Generic filler sentences
 FILLERS = [
     "The history of pottery spans thousands of years.",
     "Marine biologists study coral reef ecosystems.",
-    "Railway engineering requires precise curvature.",
-    "The periodic table arranges elements by number.",
-    "Clouds are classified as cumulus and stratus.",
-    "Beekeeping traditions differ between continents.",
-    "The Great Wall was built over many dynasties.",
-    "Thermodynamics governs heat transfer.",
-    "Impressionist painters captured fleeting light.",
-    "Volcanic activity is tracked with seismographs.",
-    "The Dewey Decimal System organizes libraries.",
-    "Irrigation evolved from canals to drip systems.",
-    "Neural networks are inspired by biological brains.",
-    "Light speed is 299,792,458 meters per second.",
-    "Classical composition follows harmonic rules.",
-    "Urban planning addresses zoning and transport.",
-    "Photosynthesis converts CO2 into glucose.",
-    "The Fibonacci sequence appears in nature.
-    "GPS uses triangulation from satellites.",
-    "Cryptography secures digital communication.",
-    "Aerodynamics explains lift on aircraft wings.",
-    "Meteorologists track pressure systems globally.",
-    "The Rosetta Stone enabled hieroglyph translation.",
-    "Bioluminescence occurs in deep ocean species.",
-    "Microprocessors revolutionized personal computing.",
-    "Tectonic plates shift gradually over millennia.",
-    "The printing press spread knowledge across Europe.",
-    "Quantum entanglement defies classical intuition.",
-    "Archaeologists excavate buried ancient settlements.",
-    "Photosensors detect minute light variations.",
+    "Railway engineering requires precise curvature calculations.",
+    "The periodic table arranges elements by atomic number.",
+    "Clouds are classified into cumulus and stratus types.",
+    "Beekeeping traditions differ significantly between continents.",
+    "The Great Wall was constructed over many successive dynasties.",
+    "Thermodynamics governs the principles of heat transfer.",
+    "Impressionist painters captured fleeting effects of light.",
+    "Volcanic activity is closely tracked with seismographs.",
+    "The Dewey Decimal System organizes library collections worldwide.",
+    "Irrigation technology evolved from canals to drip systems.",
+    "Neural networks are directly inspired by biological brains.",
+    "Light speed in vacuum is 299,792,458 meters per second.",
+    "Classical composition generally follows established harmonic rules.",
+    "Urban planning must address zoning and public transport.",
+    "Photosynthesis converts carbon dioxide into glucose and oxygen.",
+    "The Fibonacci sequence appears frequently throughout nature.",
+    "GPS navigation uses triangulation from orbiting satellites.",
+    "Cryptography secures modern digital communications against eavesdropping.",
 ]
 
-# Decoy sentences with codes that use similar patterns
-DECOY_TEMPLATES = [
-    "The transaction was logged with code TX-{code}.",
-    "The batch identifier is BC-{code}.",
-    "Session recorded under SC-{code}.",
-    "Access granted via AC-{code}.",
-    "Error logged as EC-{code}.",
-    "Debug trace shows DC-{code}.",
-    "Network packet tagged NC-{code}.",
-    "User authenticated with UC-{code}.",
-    "System heartbeat code SY-{code}.",
-    "Database query ID DB-{code}.",
-]
+# Entities that appear in MULTIPLE sentences — no single entity is unique
+NAMES = ["Alice", "Bob", "Carol", "David", "Eve", "Frank", "Grace", "Heidi"]
+ITEMS = ["bicycle", "laptop", "watch", "camera", "guitar", "sneakers", "backpack", "headphones"]
+PLACES = ["downtown shop", "uptown store", "westside mall", "eastside market", "riverside plaza"]
+
+
+def _make_entity_distractor(person: str, item: str, place: str, used: set) -> str:
+    """Create a distractor sentence sharing 1-2 entities with the target but not all 3."""
+    templates = [
+        "{person} visited the {place} last Tuesday to browse items.",
+        "The {place} sells various products including {item}s and accessories.",
+        "{person} enjoys using their {item} during weekend activities.",
+        "A customer purchased a {item} from the {place} earlier this month.",
+        "{person} recommended the {place} to friends and family members.",
+        "The {place} had a promotional sale on {item}s last holiday season.",
+        "{person} previously owned a different {item} before upgrading.",
+        "Shoppers at the {place} often look for quality {item}s.",
+    ]
+    # Pick a template and substitute with random entities (may overlap)
+    tmpl = random.choice(templates)
+    p = random.choice(NAMES)
+    it = random.choice(ITEMS)
+    pl = random.choice(PLACES)
+    sent = tmpl.format(person=p, item=it, place=pl)
+    # Ensure it shares at least one entity with the target (person, item, place)
+    # but not all three (otherwise it's a duplicate target)
+    if p == person and it == item and pl == place:
+        # Swap one entity to avoid being identical to target
+        swap = random.choice(["person", "item", "place"])
+        if swap == "person":
+            p = random.choice([n for n in NAMES if n != person])
+        elif swap == "item":
+            it = random.choice([i for i in ITEMS if i != item])
+        else:
+            pl = random.choice([pl for pl in PLACES if pl != place])
+        sent = tmpl.format(person=p, item=it, place=pl)
+    return sent
 
 
-def _make_haystack(n: int, num_decoys: int = 15) -> str:
-    """Generate n sentences of filler text with decoy codes scattered throughout."""
+def _make_haystack(n: int, target_person: str, target_item: str, target_place: str, num_distractors: int = 40) -> str:
+    """Generate n sentences with entity-overlap distractors scattered throughout."""
     sents = []
-    for i in range(n):
-        if random.random() < (num_decoys / n) and num_decoys > 0:
-            # Insert a decoy sentence instead of filler
-            template = random.choice(DECOY_TEMPLATES)
-            code = f"{random.randint(1000, 9999)}"
-            sents.append(template.format(code=code))
-            num_decoys -= 1
-        else:
-            sents.append(random.choice(FILLERS) + f" [{i+1}].")
+    
+    # Add distractor sentences that share entities
+    for _ in range(num_distractors):
+        sents.append(_make_entity_distractor(target_person, target_item, target_place, set()))
+    
+    # Fill remaining with generic fillers
+    while len(sents) < n:
+        sents.append(random.choice(FILLERS))
+    
     random.shuffle(sents)
     return " ".join(sents)
 
@@ -98,7 +114,7 @@ def run_needle_in_haystack(
     out_dir: str,
     depths: List[float] = None,
 ) -> Dict[str, Any]:
-    """Run needle-in-haystack experiment with decoy codes."""
+    """Run needle-in-haystack with entity-overlap distractors."""
     ensure_dir(out_dir)
 
     if depths is None:
@@ -111,27 +127,37 @@ def run_needle_in_haystack(
         logger.info(f"[NEEDLE] Depth {depth:.1%}")
         preds = []
         for i in tqdm(range(num_examples), desc=f"Needle {depth:.1%}", leave=False):
-            filler = _make_haystack(num_sentences, num_decoys=15)
-            code = f"{random.randint(10000, 99999)}"
-            # Target uses a different prefix than decoys to make it distinguishable
-            # but the model must find it among many codes
-            needle = f"The classified identifier is CL-{code}."
+            # Choose target entities
+            person = random.choice(NAMES)
+            item = random.choice(ITEMS)
+            place = random.choice(PLACES)
+            price = random.randint(100, 999)
+            
+            # Build haystack with entity-overlap distractors
+            filler = _make_haystack(num_sentences, person, item, place, num_distractors=40)
+            
+            # Target sentence (the needle)
+            needle = f"{person} purchased a {item} from the {place} for ${price}."
             text = _insert_needle(filler, needle, depth)
+            
+            # Question forces the model to find the RIGHT combination, not just any mention
             prompt = (
-                f"Read the text carefully. Multiple codes appear, but only one is "
-                f"the classified identifier. Find it.\n\n{text}\n\n"
-                f"What is the classified identifier? Answer with only the full code (including CL- prefix)."
+                f"Read the passage carefully. {person} is mentioned several times, "
+                f"and the {place} is mentioned several times, and {item}s are mentioned several times. "
+                f"Find the specific sentence that says how much {person} paid for a {item} at the {place}. "
+                f"Answer with only the dollar amount (no $ sign, no words)."
             )
+            
             ans = generate_text(
                 [{"role": "user", "content": prompt}],
                 model_name=model_name,
-                max_new_tokens=20,
+                max_new_tokens=10,
             )
-            correct = exact_match_score(ans, f"CL-{code}")
+            correct = exact_match_score(ans, str(price))
             preds.append({
                 "model_answer": ans,
                 "correct": correct,
-                "secret": f"CL-{code}",
+                "expected": price,
                 "depth": depth,
             })