mirage-benchmark 1.0.4__py3-none-any.whl

mirage/pipeline/deduplication.py

@@ -0,0 +1,491 @@
+
+import json
+import logging
+import sys
+import os
+from tqdm import tqdm
+from sentence_transformers import SentenceTransformer, util
+import torch
+import numpy as np
+from typing import List, Dict, Set, Tuple
+from collections import defaultdict
+
+# Import helper modules
+from mirage.embeddings.models import get_best_embedding_model
+from mirage.embeddings.rerankers_text import LLMReranker
+from mirage.core.llm import setup_logging, call_vlm_with_multiple_images
+from mirage.core.prompts import PROMPTS
+
+# Try to load configuration from config.yaml
+try:
+    from mirage.core.config import load_config
+    _cfg = load_config()
+    _dedup = _cfg.get('deduplication', {})
+    
+    QUESTION_SIMILARITY_THRESHOLD = _dedup.get('question_similarity_threshold', 0.75)
+    _ans_sim = _dedup.get('answer_similarity', {})
+    ANSWER_SIMILARITY_HIGH = _ans_sim.get('high', 0.95)
+    ANSWER_SIMILARITY_MEDIUM = _ans_sim.get('medium', 0.85)
+    ANSWER_SIMILARITY_LOW = _ans_sim.get('low', 0.70)
+    MIN_COMMUNITY_SIZE = _dedup.get('min_community_size', 2)
+    # α parameter: weight for semantic similarity vs chunk lineage (Eq. 10 in manuscript)
+    ALPHA = _dedup.get('alpha', 0.6)
+    print(f"✅ Deduplication config loaded: α={ALPHA}, question_threshold={QUESTION_SIMILARITY_THRESHOLD}")
+except ImportError:
+    print("⚠️ config_loader not available, using default deduplication configuration")
+    QUESTION_SIMILARITY_THRESHOLD = 0.75
+    ANSWER_SIMILARITY_HIGH = 0.95
+    ANSWER_SIMILARITY_MEDIUM = 0.85
+    ANSWER_SIMILARITY_LOW = 0.70
+    MIN_COMMUNITY_SIZE = 2
+    ALPHA = 0.6  # Default α for Eq. 10: α * semantic_sim + (1-α) * jaccard
+
+# Configuration
+OUTPUT_DIR = "output"
+INPUT_FILE = os.path.join(OUTPUT_DIR, "qa_multihop_pass.json")
+OUTPUT_FILE = os.path.join(OUTPUT_DIR, "qa_dataset_deduplicated.json")
+
+def load_dataset(filepath):
+    print(f"📂 Loading dataset from {filepath}...")
+    with open(filepath, 'r', encoding='utf-8') as f:
+        return json.load(f)
+
+def save_dataset(data, filepath):
+    print(f"💾 Saving {len(data)} items to {filepath}...")
+    with open(filepath, 'w', encoding='utf-8') as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+
+def compute_chunk_overlap(qa1: Dict, qa2: Dict) -> float:
+    """
+    Compute Jaccard similarity based on chunk lineage.
+    Returns overlap ratio [0, 1] of chunks used to generate each QA.
+    """
+    # Extract chunk identifiers from chunks_added (list of dicts with file_name and chunk_id)
+    def extract_chunk_ids(qa):
+        ids = {qa.get('chunk_id', -1)}
+        for chunk in qa.get('chunks_added', []):
+            if isinstance(chunk, dict):
+                ids.add((chunk.get('file_name', ''), chunk.get('chunk_id', '')))
+            else:
+                ids.add(chunk)
+        return ids
+    
+    chunks1 = extract_chunk_ids(qa1)
+    chunks2 = extract_chunk_ids(qa2)
+    
+    # Remove invalid IDs
+    chunks1.discard(-1)
+    chunks2.discard(-1)
+    
+    if not chunks1 or not chunks2:
+        return 0.0
+    
+    intersection = len(chunks1 & chunks2)
+    union = len(chunks1 | chunks2)
+    
+    return intersection / union if union > 0 else 0.0
+
+def select_best_qa(cluster_items: List[Dict]) -> Dict:
+    """
+    Select the best QA from exact duplicates based on quality metrics.
+    Prioritizes: relevance score > difficulty score > answer length
+    """
+    best = cluster_items[0]
+    best_score = (
+        float(best.get('relevance_score', 0)),
+        float(best.get('difficulty_score', 0)),
+        len(best.get('answer', ''))
+    )
+    
+    for item in cluster_items[1:]:
+        score = (
+            float(item.get('relevance_score', 0)),
+            float(item.get('difficulty_score', 0)),
+            len(item.get('answer', ''))
+        )
+        if score > best_score:
+            best = item
+            best_score = score
+    
+    return best
+
+def reorganize_qa_packs(merged_items: List[Dict], base_metadata: Dict,
+                        expert_persona: str,
+                        domain: str) -> List[Dict]:
+    """
+    Reorganize merged QA pairs into balanced question-answer packs.
+    Groups related questions together while keeping packs balanced.
+    
+    Args:
+        merged_items: List of merged QA dicts with 'question' and 'answer' keys
+        base_metadata: Metadata to propagate to reorganized items
+        expert_persona: Expert role for domain-specific organization
+        domain: Domain context for organization
+    
+    Returns:
+        List of reorganized QA dicts
+    """
+    if len(merged_items) <= 1:
+        return merged_items
+    
+    # Prepare merged questions and answers for the prompt
+    merged_questions = "\n".join([
+        f"{i+1}. {item['question']}" 
+        for i, item in enumerate(merged_items)
+    ])
+    merged_answers = "\n".join([
+        f"{i+1}. {item['answer']}" 
+        for i, item in enumerate(merged_items)
+    ])
+    
+    prompt_template = PROMPTS.get("deduplication_reorganize", "")
+    if not prompt_template:
+        logging.warning("deduplication_reorganize prompt not found, skipping reorganization.")
+        return merged_items
+    
+    # Format prompt with domain and expert role
+    formatted_prompt = prompt_template.format(expert_persona=expert_persona, domain=domain)
+    prompt = f"{formatted_prompt}\n\nInput:\nMerged Questions:\n{merged_questions}\n\nMerged Answers:\n{merged_answers}"
+    
+    try:
+        response = call_vlm_with_multiple_images(prompt, [])
+        reorganized = parse_reorganized_packs(response, base_metadata)
+        
+        if not reorganized:
+            logging.warning("LLM returned empty reorganization, keeping merged items.")
+            return merged_items
+        
+        return reorganized
+        
+    except Exception as e:
+        logging.error(f"Error in LLM reorganization: {e}")
+        return merged_items
+
+def parse_reorganized_packs(response_text: str, base_metadata: Dict) -> List[Dict]:
+    """
+    Parse the LLM response containing reorganized QA packs.
+    
+    Returns:
+        List of reorganized QA dicts
+    """
+    tuple_delimiter = PROMPTS.get("DEFAULT_TUPLE_DELIMITER", "<|#|>")
+    completion_delimiter = PROMPTS.get("DEFAULT_COMPLETION_DELIMITER", "<|#|>END<|#|>")
+    
+    qa_packs = []
+    
+    try:
+        # Remove completion delimiter if present
+        if completion_delimiter in response_text:
+            response_text = response_text.split(completion_delimiter)[0].strip()
+        
+        # Remove START delimiter if present
+        start_delimiter = tuple_delimiter + "START" + tuple_delimiter
+        if start_delimiter in response_text:
+            response_text = response_text.split(start_delimiter, 1)[-1].strip()
+        
+        # Split by NEXT delimiter to get individual packs
+        next_delimiter = tuple_delimiter + "NEXT" + tuple_delimiter
+        pack_texts = response_text.split(next_delimiter)
+        
+        for pack_text in pack_texts:
+            pack_text = pack_text.strip()
+            if not pack_text:
+                continue
+            
+            # Parse Question and Answer from the pack
+            # Expected format: Question<|#|><questions><|#|>Answer<|#|><answer>
+            if "Question" + tuple_delimiter in pack_text:
+                parts = pack_text.split(tuple_delimiter)
+                
+                question = None
+                answer = None
+                
+                for i, part in enumerate(parts):
+                    if part.lower() == "question" and i + 1 < len(parts):
+                        question = parts[i + 1].strip()
+                    elif part.lower() == "answer" and i + 1 < len(parts):
+                        answer = parts[i + 1].strip()
+                
+                if question and answer:
+                    new_item = base_metadata.copy()
+                    new_item["question"] = question
+                    new_item["answer"] = answer
+                    new_item["reorganized"] = True
+                    qa_packs.append(new_item)
+        
+        return qa_packs
+        
+    except Exception as e:
+        logging.error(f"Error parsing reorganized packs: {e}")
+        return []
+
+def hierarchical_clustering(
+    data: List[Dict],
+    question_embeddings: torch.Tensor,
+    answer_embeddings: torch.Tensor
+) -> List[List[int]]:
+    """
+    Two-stage hierarchical clustering:
+    1. Cluster questions by semantic similarity (topic/intent)
+    2. Within each question cluster, sub-cluster answers by similarity
+    
+    Returns: List of clusters (each cluster is a list of indices)
+    """
+    print(f"\n🔍 Stage 1: Clustering questions by topic (threshold: {QUESTION_SIMILARITY_THRESHOLD})...")
+    
+    # Stage 1: Cluster questions to group by topic/intent
+    question_clusters = util.community_detection(
+        question_embeddings,
+        threshold=QUESTION_SIMILARITY_THRESHOLD,
+        min_community_size=1
+    )
+    
+    print(f"✅ Found {len(question_clusters)} question-based topic groups")
+    
+    # Stage 2: Within each question cluster, sub-cluster by answer similarity
+    final_clusters = []
+    singleton_count = 0
+    
+    print(f"\n🔍 Stage 2: Sub-clustering answers within each topic group...")
+    
+    for q_cluster in tqdm(question_clusters, desc="Processing question clusters"):
+        if len(q_cluster) == 1:
+            singleton_count += 1
+            continue  # No duplicates possible
+        
+        # Extract answer embeddings for this question cluster
+        q_cluster_list = list(q_cluster)
+        q_cluster_answer_embs = answer_embeddings[q_cluster_list]
+        
+        # Check for chunk overlap to prioritize merging
+        # Build chunk overlap matrix
+        chunk_overlap_matrix = np.zeros((len(q_cluster_list), len(q_cluster_list)))
+        for i, idx_i in enumerate(q_cluster_list):
+            for j, idx_j in enumerate(q_cluster_list):
+                if i != j:
+                    chunk_overlap_matrix[i, j] = compute_chunk_overlap(data[idx_i], data[idx_j])
+        
+        # Compute answer similarity matrix
+        answer_sim_matrix = util.pytorch_cos_sim(q_cluster_answer_embs, q_cluster_answer_embs)
+        
+        # Combined similarity per Eq. 10: α * cos(e_ai, e_aj) + (1-α) * J(C^s_i, C^s_j)
+        # α weights semantic similarity; (1-α) weights chunk lineage Jaccard overlap
+        combined_sim = ALPHA * answer_sim_matrix.cpu().numpy() + (1 - ALPHA) * chunk_overlap_matrix
+        
+        # Find high-similarity pairs and group them
+        visited = set()
+        for i in range(len(q_cluster_list)):
+            if i in visited:
+                continue
+            
+            cluster = [q_cluster_list[i]]
+            visited.add(i)
+            
+            for j in range(i + 1, len(q_cluster_list)):
+                if j not in visited and combined_sim[i, j] >= ANSWER_SIMILARITY_LOW:
+                    cluster.append(q_cluster_list[j])
+                    visited.add(j)
+            
+            if len(cluster) >= MIN_COMMUNITY_SIZE:
+                final_clusters.append(cluster)
+    
+    print(f"✅ Found {len(final_clusters)} answer clusters requiring merge")
+    print(f"ℹ️  {singleton_count} singletons (unique QAs)")
+    
+    return final_clusters
+
+def process_cluster_by_similarity(
+    cluster_items: List[Dict],
+    cluster_indices: List[int],
+    answer_embeddings: torch.Tensor,
+    llm_merger: LLMReranker,
+    expert_persona: str,
+    domain: str,
+    enable_reorganization: bool = True
+) -> List[Dict]:
+    """
+    Process a cluster with stratified handling based on answer similarity.
+    - High similarity (>0.95): Exact duplicates → select best
+    - Medium similarity (0.85-0.95): Partial overlap → LLM merge → reorganize
+    - Low similarity (0.70-0.85): Related → LLM evaluate → reorganize
+    
+    After LLM merging, optionally reorganizes into balanced QA packs.
+    
+    Args:
+        cluster_items: List of QA dicts in the cluster
+        cluster_indices: Original indices of items in the cluster
+        answer_embeddings: Tensor of answer embeddings
+        llm_merger: LLMReranker instance for merging
+        enable_reorganization: Whether to reorganize after merging
+        expert_persona: Expert role for domain-specific handling
+        domain: Domain context for handling
+    """
+    if len(cluster_items) < 2:
+        return cluster_items
+    
+    # Compute pairwise answer similarities within cluster
+    cluster_answer_embs = answer_embeddings[cluster_indices]
+    sim_matrix = util.pytorch_cos_sim(cluster_answer_embs, cluster_answer_embs).cpu().numpy()
+    
+    # Get max similarity (excluding diagonal)
+    np.fill_diagonal(sim_matrix, 0)
+    max_similarity = np.max(sim_matrix)
+    avg_similarity = np.mean(sim_matrix[np.triu_indices_from(sim_matrix, k=1)])
+    
+    # Base metadata for propagation
+    base_metadata = cluster_items[0].copy()
+    base_metadata.pop('question', None)
+    base_metadata.pop('answer', None)
+    base_metadata['merged_from_count'] = len(cluster_items)
+    
+    # Stratified handling
+    if max_similarity >= ANSWER_SIMILARITY_HIGH:
+        # Tier 1: Exact duplicates - just pick the best one (no reorganization needed)
+        logging.info(f"Cluster of {len(cluster_items)} with max_sim={max_similarity:.3f}: Selecting best (exact duplicates)")
+        best = select_best_qa(cluster_items)
+        best['dedup_method'] = 'select_best'
+        best['merged_from_count'] = len(cluster_items)
+        return [best]
+    
+    elif avg_similarity >= ANSWER_SIMILARITY_MEDIUM:
+        # Tier 2: High overlap - merge with LLM, then reorganize
+        logging.info(f"Cluster of {len(cluster_items)} with avg_sim={avg_similarity:.3f}: LLM merge (high overlap)")
+        merged = llm_merger.deduplicate_and_merge(cluster_items)
+        
+        # Reorganize into balanced packs if enabled and multiple items
+        if enable_reorganization and len(merged) > 1:
+            logging.info(f"Reorganizing {len(merged)} merged items into balanced packs")
+            merged = reorganize_qa_packs(merged, base_metadata, expert_persona=expert_persona, domain=domain)
+        
+        for item in merged:
+            item['dedup_method'] = 'llm_merge_high'
+        return merged
+    
+    else:
+        # Tier 3: Related but potentially distinct - let LLM decide, then reorganize
+        logging.info(f"Cluster of {len(cluster_items)} with avg_sim={avg_similarity:.3f}: LLM evaluate (medium overlap)")
+        merged = llm_merger.deduplicate_and_merge(cluster_items)
+        
+        # Reorganize into balanced packs if enabled and multiple items
+        if enable_reorganization and len(merged) > 1:
+            logging.info(f"Reorganizing {len(merged)} merged items into balanced packs")
+            merged = reorganize_qa_packs(merged, base_metadata, expert_persona=expert_persona, domain=domain)
+        
+        for item in merged:
+            item['dedup_method'] = 'llm_merge_medium'
+        return merged
+
+def deduplicate_dataset():
+    setup_logging()
+    
+    # 1. Load Data
+    if not os.path.exists(INPUT_FILE):
+        print(f"❌ Input file {INPUT_FILE} not found.")
+        return
+
+    data = load_dataset(INPUT_FILE)
+    if not data:
+        print("⚠️ Dataset is empty.")
+        return
+
+    print(f"\n{'='*80}")
+    print("🎯 HIERARCHICAL DEDUPLICATION: Questions → Answers")
+    print(f"{'='*80}\n")
+
+    # 2. Prepare separate embeddings for questions and answers
+    print("⚙️  Preparing text for embedding...")
+    questions = [item['question'] for item in data]
+    answers = [item['answer'] for item in data]
+
+    # 3. Load Embedding Model & Embed
+    model_name = get_best_embedding_model()
+    print(f"🤖 Loading embedding model: {model_name}")
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    embedder = SentenceTransformer(model_name, device=device)
+    
+    print(f"📊 Generating question embeddings for {len(questions)} QA pairs...")
+    question_embeddings = embedder.encode(questions, convert_to_tensor=True, show_progress_bar=True)
+    
+    print(f"📊 Generating answer embeddings for {len(answers)} QA pairs...")
+    answer_embeddings = embedder.encode(answers, convert_to_tensor=True, show_progress_bar=True)
+
+    # 4. Hierarchical Clustering: Questions first, then Answers
+    clusters = hierarchical_clustering(data, question_embeddings, answer_embeddings)
+    
+    # 5. Track processed items
+    clustered_indices = set()
+    for cluster in clusters:
+        clustered_indices.update(cluster)
+
+    # Initialize LLM merger
+    llm_merger = LLMReranker()
+    final_dataset = []
+    
+    # Statistics
+    stats = {
+        'original': len(data),
+        'singletons': 0,
+        'clusters_processed': 0,
+        'items_in_clusters': 0,
+        'exact_duplicates': 0,
+        'llm_merges': 0,
+        'reorganized_packs': 0
+    }
+
+    # Add singletons (items not in any cluster)
+    for i in range(len(data)):
+        if i not in clustered_indices:
+            final_dataset.append(data[i])
+            stats['singletons'] += 1
+
+    print(f"\nℹ️  Added {stats['singletons']} unique (singleton) items.")
+
+    # 6. Process clusters with stratified handling
+    print(f"\n🔄 Processing {len(clusters)} clusters with stratified merge strategy...")
+    
+    for cluster in tqdm(clusters, desc="Merging clusters"):
+        cluster_items = [data[idx] for idx in cluster]
+        stats['clusters_processed'] += 1
+        stats['items_in_clusters'] += len(cluster_items)
+        
+        # Process with stratified approach
+        merged_items = process_cluster_by_similarity(
+            cluster_items, 
+            cluster, 
+            answer_embeddings,
+            llm_merger
+        )
+        
+        # Track merge type
+        if merged_items and merged_items[0].get('dedup_method') == 'select_best':
+            stats['exact_duplicates'] += 1
+        else:
+            stats['llm_merges'] += 1
+        
+        # Track reorganized packs
+        reorganized_count = sum(1 for item in merged_items if item.get('reorganized', False))
+        if reorganized_count > 0:
+            stats['reorganized_packs'] += reorganized_count
+        
+        final_dataset.extend(merged_items)
+
+    # 7. Save Results
+    print("\n" + "="*80)
+    print("📊 HIERARCHICAL DEDUPLICATION SUMMARY")
+    print("="*80)
+    print(f"Original count:           {stats['original']}")
+    print(f"Final count:              {len(final_dataset)}")
+    print(f"Reduction:                {stats['original'] - len(final_dataset)} ({100*(stats['original'] - len(final_dataset))/stats['original']:.1f}%)")
+    print(f"---")
+    print(f"Singleton items:          {stats['singletons']}")
+    print(f"Clusters processed:       {stats['clusters_processed']}")
+    print(f"Items in clusters:        {stats['items_in_clusters']}")
+    print(f"Exact duplicates removed: {stats['exact_duplicates']}")
+    print(f"LLM merges performed:     {stats['llm_merges']}")
+    print(f"Reorganized QA packs:     {stats['reorganized_packs']}")
+    print("="*80)
+    
+    save_dataset(final_dataset, OUTPUT_FILE)
+
+if __name__ == "__main__":
+    deduplicate_dataset()