mirage-benchmark 1.0.4__py3-none-any.whl

mirage/embeddings/rerankers_multimodal.py

@@ -0,0 +1,766 @@
+import torch
+import re
+from abc import ABC, abstractmethod
+from typing import List, Optional, Dict, Tuple
+from pathlib import Path
+from PIL import Image
+
+class BaseReranker(ABC):
+    """Abstract base class for image-based rerankers"""
+    
+    @abstractmethod
+    def rerank(self, query: str, image_paths: List[str], top_k: int = 10) -> List[int]:
+        pass
+
+class ChunkReranker(ABC):
+    """Abstract base class for chunk-based rerankers (text + optional images)"""
+    
+    @abstractmethod
+    def rerank(self, query: str, chunks: List[Dict[str, str]], top_k: int = 1) -> List[Tuple[int, float, Dict[str, str]]]:
+        """
+        Rerank chunks based on query relevance
+        
+        Args:
+            query: User query string
+            chunks: List of dicts with 'text' and optional 'image_path' keys
+            top_k: Number of top chunks to return
+        
+        Returns:
+            List of tuples: (original_index, relevance_score, chunk_dict)
+        """
+        pass
+
+class MMR5Reranker(BaseReranker):
+    """MM-R5: MultiModal Reasoning-Enhanced ReRanker"""
+    
+    SYSTEM_PROMPT = (
+        "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant "
+        "first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning "
+        "process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
+        "<think> reasoning process here </think><answer> answer here </answer>"
+    )
+    
+    QUESTION_TEMPLATE = (
+        "Please rank the following images according to their relevance to the question. "
+        "Provide your response in the format: <think>your reasoning process here</think><answer>[image_id_1, image_id_2, ...]</answer> "
+        "where the numbers in the list represent the ranking order of images'id from most to least relevant. "
+        "Before outputting the answer, you need to analyze each image and provide your analysis process."
+        "For example: <think>Image 1 shows the most relevant content because...</think><answer>[id_most_relevant, id_second_relevant, ...]</answer>"
+        "\nThe question is: {Question}"
+        "\n\nThere are {image_num} images, id from 1 to {image_num_end}, Image ID to image mapping:\n"
+    )
+    
+    def __init__(self, model_name: str = "i2vec/MM-R5"):
+        print(f"Loading MM-R5: {model_name}")
+        
+        # Try official MM-R5 package first
+        try:
+            from reranker import QueryReranker  # type: ignore
+            self.reranker = QueryReranker(model_name)
+            self.use_official_package = True
+            print(f"✅ MM-R5 loaded via official package")
+        except ImportError:
+            # Fallback to direct implementation using official code
+            from transformers import Qwen2_5_VLProcessor, Qwen2_5_VLForConditionalGeneration
+            from qwen_vl_utils import process_vision_info
+            
+            attn_kwargs = {}
+            if self._flash_attn_available():
+                attn_kwargs["attn_implementation"] = "flash_attention_2"
+            else:
+                print("⚠️  flash_attn not available. Falling back to default attention implementation.")
+            
+            # Load without device_map="auto" to avoid meta tensor issues in parallel processing
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model_name,
+                torch_dtype=torch.bfloat16,
+                low_cpu_mem_usage=True,
+                **attn_kwargs,
+            ).to(self.device).eval()
+            
+            self.processor = Qwen2_5_VLProcessor.from_pretrained(model_name)
+            self.use_official_package = False
+            print(f"✅ MM-R5 loaded via direct implementation")
+    
+    def rerank(self, query: str, image_paths: List[str], top_k: int = 10) -> List[int]:
+        """Rerank images based on query relevance"""
+        import re
+        
+        if self.use_official_package:
+            # Use official package API
+            predicted_order = self.reranker.rerank(query, image_paths)
+        else:
+            # Use direct implementation following official code
+            from qwen_vl_utils import process_vision_info
+            
+            device = self.model.device
+            
+            messages = [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": self.SYSTEM_PROMPT,
+                        },
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": self.QUESTION_TEMPLATE.format(
+                                Question=query,
+                                image_num=len(image_paths),
+                                image_num_end=len(image_paths)
+                            ),
+                        },
+                    ],
+                },
+            ]
+            
+            # Add images to messages
+            for i, image_path in enumerate(image_paths):
+                messages[-1]["content"].extend(
+                    [
+                        {"type": "text", "text": f"\nImage {i+1}: "},
+                        {"type": "image", "image": image_path},
+                    ]
+                )
+                
+            text = self.processor.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+            
+            image_inputs, video_inputs = process_vision_info(messages)
+            
+            inputs = self.processor(
+                text=[text],
+                images=image_inputs,
+                videos=video_inputs,
+                padding=True,
+                return_tensors="pt",
+            )
+            
+            inputs = inputs.to(device)
+            
+            generated_ids = self.model.generate(
+                **inputs,
+                do_sample=True,
+                temperature=0.3,
+                max_new_tokens=8192,
+                use_cache=True,
+            )
+
+            generated_ids_trimmed = [
+                out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+            ]
+            
+            output_text = self.processor.batch_decode(
+                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+            )[0]
+            
+            # Parse output results
+            match = re.search(r'<answer>\[(.*?)\]</answer>', output_text)
+            
+            if match:
+                try:
+                    tmp_predicted_order = []
+                    predicted_order = [int(x) - 1 for x in match.group(1).strip().split(',') if x.strip()]
+                    
+                    for idx in predicted_order:
+                        if 0 <= idx < len(image_paths):
+                            tmp_predicted_order.append(idx)
+                            
+                    predicted_order = tmp_predicted_order
+                    
+                    # Handle missing indices
+                    if len(set(predicted_order)) < len(image_paths):
+                        missing_ids = set(range(len(image_paths))) - set(predicted_order)
+                        predicted_order.extend(sorted(list(missing_ids)))
+                        
+                except Exception as e:
+                    predicted_order = [i for i in range(len(image_paths))]
+                    print(f"⚠️  Parsing error: {str(e)}, output text: {output_text[:200]}...")
+            else:
+                predicted_order = [i for i in range(len(image_paths))]
+                print(f"⚠️  Could not parse ranking from output: {output_text[:200]}...")
+        
+        return predicted_order[:top_k]
+    
+    @staticmethod
+    def _flash_attn_available() -> bool:
+        try:
+            import flash_attn  # noqa: F401
+            return True
+        except Exception:
+            return False
+
+class Florence2Reranker(BaseReranker):
+    """Florence-2-large for visual document reranking"""
+    
+    def __init__(self, model_name: str = "microsoft/Florence-2-large"):
+        from transformers import AutoProcessor, AutoModelForCausalLM
+        try:
+            from transformers import BitsAndBytesConfig
+        except ImportError:
+            BitsAndBytesConfig = None
+        
+        print(f"Loading Florence-2: {model_name}")
+        
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        
+        # Note: When using quantization, device_map is needed for BitsAndBytes
+        # When not using quantization, load explicitly to avoid meta tensor issues
+        if self.device == "cuda":
+            try:
+                quantization_config = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch.float16,
+                    bnb_4bit_use_double_quant=True,
+                )
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    torch_dtype=self.torch_dtype,
+                    trust_remote_code=True,
+                    quantization_config=quantization_config,
+                    device_map="auto"  # Required for BitsAndBytes quantization
+                ).eval()
+            except Exception as e:
+                print(f"⚠️  Quantization failed ({e}), loading without quantization...")
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    torch_dtype=self.torch_dtype,
+                    trust_remote_code=True,
+                    low_cpu_mem_usage=True,
+                ).to(self.device).eval()
+        else:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=self.torch_dtype,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+            ).to(self.device).eval()
+        
+        self.processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+        print(f"✅ Florence-2 loaded on {self.device}")
+    
+    def _score_image(self, query: str, image_path: str) -> float:
+        """Score a single image based on query relevance"""
+        try:
+            if not Path(image_path).exists():
+                return 0.0
+            
+            image = Image.open(image_path).convert('RGB')
+            
+            # Use caption task to understand image
+            prompt = "<CAPTION>"
+            inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(self.device)
+            
+            with torch.no_grad():
+                generated_ids = self.model.generate(
+                    input_ids=inputs["input_ids"],
+                    pixel_values=inputs["pixel_values"],
+                    max_new_tokens=1024,
+                    num_beams=3,
+                    do_sample=False,
+                    use_cache=False,
+                )
+            
+            generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+            parsed_answer = self.processor.post_process_generation(
+                generated_text, task="<CAPTION>", image_size=(image.width, image.height)
+            )
+            
+            caption = parsed_answer.get('<CAPTION>', '')
+            
+            # Simple relevance score based on query keyword overlap
+            query_words = set(query.lower().split())
+            caption_words = set(caption.lower().split())
+            overlap = len(query_words & caption_words)
+            score = overlap / max(len(query_words), 1)
+            
+            return score
+            
+        except Exception as e:
+            print(f"⚠️  Scoring failed for {image_path}: {e}")
+            return 0.0
+    
+    def rerank(self, query: str, image_paths: List[str], top_k: int = 10) -> List[int]:
+        """Rerank images based on query relevance using Florence-2"""
+        try:
+            scores = []
+            for img_path in image_paths:
+                score = self._score_image(query, img_path)
+                scores.append(score)
+            
+            # Get top-k indices sorted by score
+            ranked_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
+            return ranked_indices[:top_k]
+            
+        except Exception as e:
+            print(f"⚠️  Reranking failed: {e}")
+            return list(range(min(top_k, len(image_paths))))
+
+class VLMReranker(ChunkReranker):
+    """VLM-based reranker using Motor Maven endpoint with multiple images"""
+    
+    def __init__(self):
+        from call_llm import call_vlm_with_multiple_images
+        from prompt import PROMPTS
+        
+        self.call_vlm_multi = call_vlm_with_multiple_images
+        self.rerank_prompt = PROMPTS["rerank_vlm"]
+        print("✅ VLM Reranker initialized")
+    
+    def rerank(self, query: str, chunks: List[Dict[str, str]], top_k: int = 1) -> List[Tuple[int, float, Dict[str, str]]]:
+        """
+        Rerank chunks based on query relevance with multiple images
+        
+        Args:
+            query: User query string
+            chunks: List of dicts with 'text' and optional 'image_path' keys
+            top_k: Number of top chunks to return (default: 1)
+        
+        Returns:
+            List of tuples: (original_index, relevance_score, chunk_dict)
+        """
+        try:
+            # Collect chunk data and images
+            chunk_data = []
+            image_paths = []
+            chunk_to_image_map = {}  # Maps chunk index to image index
+            
+            for i, chunk in enumerate(chunks):
+                chunk_info = {
+                    'index': i,
+                    'text': chunk.get('text', ''),
+                    'image_path': chunk.get('image_path', None),
+                    'has_image': False
+                }
+                
+                # Track images
+                if chunk_info['image_path'] and Path(chunk_info['image_path']).exists():
+                    chunk_info['has_image'] = True
+                    chunk_to_image_map[i] = len(image_paths)
+                    image_paths.append(chunk_info['image_path'])
+                
+                chunk_data.append(chunk_info)
+            
+            # Build structured prompt with explicit chunk boundaries
+            formatted_chunks = []
+            for chunk_info in chunk_data:
+                i = chunk_info['index'] + 1  # 1-indexed for display
+                chunk_lines = [f"<CHUNK_START id={i}>"]
+                chunk_lines.append(chunk_info['text'])
+                
+                if chunk_info['has_image']:
+                    img_idx = chunk_to_image_map[chunk_info['index']] + 1
+                    chunk_lines.append(f"<IMAGE_START id={img_idx} relates_to_chunk={i}>")
+                    chunk_lines.append(f"[Image {img_idx} displayed here]")
+                    chunk_lines.append(f"<IMAGE_END id={img_idx}>")
+                
+                chunk_lines.append(f"<CHUNK_END id={i}>")
+                formatted_chunks.append("\n".join(chunk_lines))
+            
+            # Build full prompt
+            full_prompt = f"""{self.rerank_prompt}
+
+Query: {query}
+
+Chunks to rank:
+
+{chr(10).join(formatted_chunks)}"""
+            
+            # Call VLM with all images
+            if not image_paths:
+                # No images - use LLM fallback instead of VLM
+                from call_llm import call_llm_simple
+                response = call_llm_simple(full_prompt)
+            else:
+                response = self.call_vlm_multi(full_prompt, image_paths)
+            
+            # Parse response to extract rankings
+            rankings = self._parse_rankings(response, chunk_data)
+            
+            # Return top-k with chunk data
+            result = []
+            for idx, score in rankings[:top_k]:
+                result.append((idx, score, chunks[idx]))
+            
+            return result
+            
+        except Exception as e:
+            print(f"⚠️  Reranking failed: {e}")
+            import traceback
+            traceback.print_exc()
+            # Return original order with default scores
+            return [(i, 1.0, chunks[i]) for i in range(min(top_k, len(chunks)))]
+    
+    def _parse_rankings(self, response: str, chunk_data: List[Dict]) -> List[Tuple[int, float]]:
+        """Parse VLM response to extract chunk rankings from structured format"""
+        rankings = []
+        num_chunks = len(chunk_data)
+        
+        # Primary pattern: <Rank X>Chunk Y (simplified format)
+        rank_pattern = r'<Rank\s+(\d+)>\s*Chunk\s+(\d+)'
+        matches = re.findall(rank_pattern, response, re.IGNORECASE)
+        
+        seen_indices = set()
+        for rank_num, chunk_num in matches:
+            idx = int(chunk_num) - 1  # Convert to 0-indexed
+            rank = int(rank_num)
+            
+            # Calculate score based on rank (higher rank = lower score)
+            # Rank 1 gets highest score (1.0), decreasing linearly
+            relevance = 1.0 - ((rank - 1) / max(num_chunks, 1))
+            
+            # Ensure valid index and no duplicates
+            if 0 <= idx < num_chunks and idx not in seen_indices:
+                rankings.append((idx, relevance))
+                seen_indices.add(idx)
+        
+        # If parsing failed or incomplete, fill remaining chunks with low scores
+        if len(rankings) < num_chunks:
+            missing = set(range(num_chunks)) - seen_indices
+            for idx in missing:
+                rankings.append((idx, 0.0))
+            print(f"⚠️  Parsed {len(seen_indices)}/{num_chunks} chunks from response")
+            if len(seen_indices) == 0:
+                # Debug: print response when parsing completely fails
+                print(f"🔍 Debug - VLM Response (first 500 chars):\n{response[:500]}")
+                print(f"🔍 Debug - VLM Response (last 500 chars):\n{response[-500:]}")
+        
+        # Sort by relevance score (highest first)
+        rankings.sort(key=lambda x: x[1], reverse=True)
+        
+        return rankings
+
+class MonoVLMReranker(ChunkReranker):
+    """MonoVLM reranker using lightonai/MonoQwen2-VL-v0.1"""
+    
+    def __init__(
+        self,
+        model_name: str = "lightonai/MonoQwen2-VL-v0.1",
+        processor_name: str = "Qwen/Qwen2-VL-2B-Instruct",
+    ):
+        print(f"Loading MonoVLM: {model_name}")
+        from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
+        
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+        
+        self.processor = AutoProcessor.from_pretrained(processor_name, trust_remote_code=True)
+        
+        # Load model without device_map="auto" to avoid meta tensor issues
+        # First load to CPU, then move to GPU explicitly
+        self.model = Qwen2VLForConditionalGeneration.from_pretrained(
+            model_name,
+            torch_dtype=self.torch_dtype,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,  # Reduces CPU memory during loading
+        )
+        
+        # Move to device and set to eval mode
+        self.model = self.model.to(self.device).eval()
+        
+        # Cache token IDs for True/False classification
+        tokenizer = getattr(self.processor, "tokenizer", None)
+        if tokenizer is None:
+            raise ValueError("Processor does not expose a tokenizer needed for MonoVLM scoring.")
+        
+        self.true_token_id = tokenizer.convert_tokens_to_ids("True")
+        self.false_token_id = tokenizer.convert_tokens_to_ids("False")
+        
+        if self.true_token_id is None or self.false_token_id is None:
+            raise ValueError("Tokenizer missing True/False tokens required for MonoVLM scoring.")
+        
+        print(f"✅ MonoVLM loaded")
+    
+    def _build_prompt(self, query: str, chunk_text: str) -> str:
+        chunk_text = chunk_text.strip() if chunk_text else "[No text provided]"
+        return (
+            "Assert the relevance of the provided document (text and/or image) to the query.\n"
+            "Respond with a single word: True if relevant, otherwise False.\n\n"
+            f"Query:\n{query}\n\nDocument:\n{chunk_text}"
+        )
+    
+    def _score_chunk(self, query: str, chunk: Dict[str, str]) -> float:
+        image = None
+        image_path = chunk.get('image_path')
+        
+        try:
+            if image_path and Path(image_path).exists():
+                with Image.open(image_path) as img:
+                    image = img.convert("RGB")
+        except Exception as img_err:
+            print(f"⚠️  Failed to load image {image_path}: {img_err}")
+            image = None
+        
+        prompt = self._build_prompt(query, chunk.get('text', ''))
+        
+        messages = [
+            {
+                "role": "user",
+                "content": (
+                    [{"type": "image", "image": image}] if image else []
+                ) + [
+                    {
+                        "type": "text",
+                        "text": prompt,
+                    }
+                ],
+            }
+        ]
+        
+        try:
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            
+            processor_kwargs = {
+                "text": text,
+                "return_tensors": "pt",
+            }
+            
+            if image is not None:
+                processor_kwargs["images"] = image
+            
+            inputs = self.processor(**processor_kwargs)
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+            
+            logits = outputs.logits[:, -1, :]
+            relevance = torch.softmax(
+                logits[:, [self.true_token_id, self.false_token_id]],
+                dim=-1
+            )
+            
+            return relevance[0, 0].item()
+        
+        except Exception as e:
+            print(f"⚠️  MonoVLM scoring failed: {e}")
+            return 0.0
+    
+    def rerank(self, query: str, chunks: List[Dict[str, str]], top_k: int = 1) -> List[Tuple[int, float, Dict[str, str]]]:
+        """
+        Rerank chunks based on query relevance using MonoVLM
+        
+        Args:
+            query: User query string
+            chunks: List of dicts with 'text' and optional 'image_path' keys
+            top_k: Number of top chunks to return (default: 1)
+        
+        Returns:
+            List of tuples: (original_index, relevance_score, chunk_dict)
+        """
+        scores = []
+        for idx, chunk in enumerate(chunks):
+            score = self._score_chunk(query, chunk)
+            scores.append((idx, score))
+        
+        scores.sort(key=lambda x: x[1], reverse=True)
+        
+        result = []
+        for idx, score in scores[:top_k]:
+            if 0 <= idx < len(chunks):
+                result.append((idx, score, chunks[idx]))
+        
+        return result
+
+class TextEmbeddingReranker(ChunkReranker):
+    """Text embedding reranker using BAAI/bge-large-en-v1.5 with image descriptions"""
+    
+    def __init__(self, model_name: str = "BAAI/bge-large-en-v1.5"):
+        from sentence_transformers import SentenceTransformer
+        
+        print(f"Loading text embedding model: {model_name}")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        self.model = SentenceTransformer(model_name, device=self.device)
+        print(f"✅ Text embedding model loaded on {self.device}")
+        
+        # For generating image descriptions (reuse VLM)
+        from call_llm import call_vlm_simple
+        from prompt import PROMPTS
+        self.call_vlm = call_vlm_simple
+        self.desc_prompt = PROMPTS.get("rerank_image_desc", "Generate a concise 100-word technical description of this image.")
+    
+    def _generate_image_description(self, image_path: str) -> str:
+        """Generate description for image using VLM"""
+        try:
+            return self.call_vlm(self.desc_prompt, image_path)
+        except Exception as e:
+            print(f"⚠️  Image description failed for {image_path}: {e}")
+            return "[Image description unavailable]"
+    
+    def rerank(self, query: str, chunks: List[Dict[str, str]], top_k: int = 1) -> List[Tuple[int, float, Dict[str, str]]]:
+        """
+        Rerank chunks based on query relevance using text embeddings
+        
+        Args:
+            query: User query string
+            chunks: List of dicts with 'text' and optional 'image_path' keys
+            top_k: Number of top chunks to return (default: 1)
+        
+        Returns:
+            List of tuples: (original_index, relevance_score, chunk_dict)
+        """
+        try:
+            # Generate image descriptions and combine with text
+            chunk_texts = []
+            for chunk in chunks:
+                text = chunk.get('text', '')
+                
+                # Add image description if image exists
+                if chunk.get('image_path') and Path(chunk['image_path']).exists():
+                    img_desc = self._generate_image_description(chunk['image_path'])
+                    combined_text = f"{text}\n[Image Description: {img_desc}]"
+                else:
+                    combined_text = text
+                
+                chunk_texts.append(combined_text)
+            
+            # Compute embeddings
+            query_embedding = self.model.encode([query], convert_to_tensor=True, normalize_embeddings=True)
+            chunk_embeddings = self.model.encode(chunk_texts, convert_to_tensor=True, normalize_embeddings=True)
+            
+            # Compute cosine similarities
+            similarities = torch.nn.functional.cosine_similarity(
+                query_embedding, chunk_embeddings, dim=1
+            )
+            
+            # Get top-k indices sorted by similarity
+            top_indices = torch.argsort(similarities, descending=True)[:top_k].cpu().tolist()
+            
+            # Build results with scores
+            results = []
+            for idx in top_indices:
+                if 0 <= idx < len(chunks):
+                    score = float(similarities[idx].cpu())
+                    results.append((idx, score, chunks[idx]))
+            
+            return results
+            
+        except Exception as e:
+            print(f"⚠️  Reranking failed: {e}")
+            return [(i, 1.0, chunks[i]) for i in range(min(top_k, len(chunks)))]
+
+if __name__ == "__main__":
+    from pathlib import Path
+    
+    print("=" * 60)
+    print("Testing Multimodal Rerankers")
+    print("=" * 60)
+    
+    # Prepare test data (text-only chunks for testing without specific images)
+    test_chunks = [
+        {
+            "text": "Machine learning models require careful hyperparameter tuning to achieve optimal performance. Common parameters include learning rate, batch size, and regularization strength.",
+            "image_path": None
+        },
+        {
+            "text": "This figure shows the characteristic relationship between model accuracy and training epochs. The data demonstrates typical learning curve behavior with diminishing returns after initial rapid improvement.",
+            "image_path": None  # Set to actual image path for multimodal testing
+        },
+        {
+            "text": "Neural network architectures vary widely in depth and complexity. Convolutional neural networks excel at image tasks, while transformers dominate natural language processing applications.",
+            "image_path": None
+        },
+        {
+            "text": "Data preprocessing is essential for model performance. Standard techniques include normalization, handling missing values, and feature engineering for tabular data.",
+            "image_path": None
+        },
+        {
+            "text": "This flowchart illustrates the machine learning pipeline from data collection through model deployment, showing how each stage contributes to the final system performance.",
+            "image_path": None  # Set to actual image path for multimodal testing
+        }
+    ]
+    
+    # Extract valid image paths for image-based rerankers
+    image_paths = [chunk['image_path'] for chunk in test_chunks if chunk.get('image_path') and Path(chunk['image_path']).exists()]
+    valid_chunks = [chunk for chunk in test_chunks if chunk['image_path'] is None or (chunk.get('image_path') and Path(chunk['image_path']).exists())]
+    query = "How do machine learning models improve with training, and what are the key stages in the ML pipeline?"
+    
+    # Test 1: VLM Reranker
+    print("\n1. Testing VLM Reranker...")
+    print("-" * 60)
+    try:
+        if valid_chunks:
+            reranker = VLMReranker()
+            print(f"Testing with {len(valid_chunks)} chunks")
+            print(f"Chunks with images: {sum(1 for c in valid_chunks if c.get('image_path'))}")
+            print(f"\nQuery: {query}\n")
+            results = reranker.rerank(query, valid_chunks, top_k=3)
+            print("\nTop 3 Reranked Chunks:")
+            for i, (orig_idx, score, chunk) in enumerate(results, 1):
+                print(f"\n{i}. Original Index: {orig_idx}, Relevance Score: {score:.3f}")
+                print(f"   Text preview: {chunk['text'][:100]}...")
+                if chunk.get('image_path'):
+                    print(f"   Has image: {Path(chunk['image_path']).name}")
+                else:
+                    print(f"   Text-only chunk")
+            print("\n✅ VLM Reranker test completed!")
+        else:
+            print("⚠️  No valid test chunks found.")
+    except Exception as e:
+        print(f"❌ VLM Reranker test failed: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    # Test 2: MonoVLM Reranker
+    print("\n2. Testing MonoVLM Reranker (lightonai/MonoQwen2-VL-v0.1)...")
+    print("-" * 60)
+    try:
+        if valid_chunks:
+            reranker = MonoVLMReranker()
+            print(f"Testing with {len(valid_chunks)} chunks")
+            print(f"Chunks with images: {sum(1 for c in valid_chunks if c.get('image_path'))}")
+            print(f"\nQuery: {query}\n")
+            results = reranker.rerank(query, valid_chunks, top_k=3)
+            print("\nTop 3 Reranked Chunks:")
+            for i, (orig_idx, score, chunk) in enumerate(results, 1):
+                print(f"\n{i}. Original Index: {orig_idx}, Relevance Score: {score:.3f}")
+                print(f"   Text preview: {chunk['text'][:100]}...")
+                if chunk.get('image_path'):
+                    print(f"   Has image: {Path(chunk['image_path']).name}")
+                else:
+                    print(f"   Text-only chunk")
+            print("\n✅ MonoVLM Reranker test completed!")
+        else:
+            print("⚠️  No valid test chunks found.")
+    except Exception as e:
+        print(f"⚠️  MonoVLM Reranker test skipped: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    # Test 3: Text Embedding Reranker (BAAI/bge-large-en-v1.5)
+    print("\n3. Testing Text Embedding Reranker (BAAI/bge-large-en-v1.5)...")
+    print("-" * 60)
+    try:
+        if valid_chunks:
+            reranker = TextEmbeddingReranker()
+            print(f"Testing with {len(valid_chunks)} chunks")
+            print(f"Chunks with images: {sum(1 for c in valid_chunks if c.get('image_path'))}")
+            print(f"\nQuery: {query}\n")
+            results = reranker.rerank(query, valid_chunks, top_k=3)
+            print("\nTop 3 Reranked Chunks:")
+            for i, (orig_idx, score, chunk) in enumerate(results, 1):
+                print(f"\n{i}. Original Index: {orig_idx}, Relevance Score: {score:.3f}")
+                print(f"   Text preview: {chunk['text'][:100]}...")
+                if chunk.get('image_path'):
+                    print(f"   Has image: {Path(chunk['image_path']).name}")
+                else:
+                    print(f"   Text-only chunk")
+            print("\n✅ Text Embedding Reranker test completed!")
+        else:
+            print("⚠️  No valid test chunks found.")
+    except Exception as e:
+        print(f"⚠️  Text Embedding Reranker test skipped: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    print("\n" + "=" * 60)