diff --git a/rag/__init__.py b/rag/__init__.py new file mode 100644 index 0000000..cd2b504 --- /dev/null +++ b/rag/__init__.py @@ -0,0 +1,8 @@ +""" +RAG (Retrieval-Augmented Generation) System for DriveMind-V +""" +from .vector_store import VectorStore +from .retriever import CLIPRetriever +from .rag_vlm import RAG_VLM + +__all__ = ['VectorStore', 'CLIPRetriever', 'RAG_VLM'] diff --git a/rag/__pycache__/__init__.cpython-310.pyc b/rag/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..8b22114 Binary files /dev/null and b/rag/__pycache__/__init__.cpython-310.pyc differ diff --git a/rag/__pycache__/rag_vlm.cpython-310.pyc b/rag/__pycache__/rag_vlm.cpython-310.pyc new file mode 100644 index 0000000..09511e6 Binary files /dev/null and b/rag/__pycache__/rag_vlm.cpython-310.pyc differ diff --git a/rag/__pycache__/retriever.cpython-310.pyc b/rag/__pycache__/retriever.cpython-310.pyc new file mode 100644 index 0000000..cd65273 Binary files /dev/null and b/rag/__pycache__/retriever.cpython-310.pyc differ diff --git a/rag/__pycache__/vector_store.cpython-310.pyc b/rag/__pycache__/vector_store.cpython-310.pyc new file mode 100644 index 0000000..97f588e Binary files /dev/null and b/rag/__pycache__/vector_store.cpython-310.pyc differ diff --git a/rag/rag_vlm.py b/rag/rag_vlm.py new file mode 100644 index 0000000..525ceb2 --- /dev/null +++ b/rag/rag_vlm.py @@ -0,0 +1,219 @@ +""" +RAG-Enhanced Vision-Language Model +Combines retrieval with generation for few-shot reasoning +""" +import os +import sys +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from typing import Union, List, Dict +from PIL import Image +from model.model_vlm import MiniMindVLM +from rag.retriever import CLIPRetriever +from rag.vector_store import VectorStore + + +class RAG_VLM: + """ + RAG-enhanced VLM for improved generation through retrieval + + Workflow: + 1. Encode query image with CLIP + 2. Retrieve k similar images from knowledge base + 3. Build few-shot prompt with retrieved examples + 4. Generate response with VLM + """ + + def __init__( + self, + vlm_model: MiniMindVLM, + tokenizer, + vector_store: VectorStore, + retriever: CLIPRetriever = None, + k: int = 3 + ): + """ + Initialize RAG-VLM + + Args: + vlm_model: MiniMindVLM instance + tokenizer: Tokenizer for VLM + vector_store: VectorStore with image knowledge base + retriever: CLIPRetriever instance (creates new one if None) + k: Number of examples to retrieve + """ + self.model = vlm_model + self.tokenizer = tokenizer + self.vector_store = vector_store + self.k = k + + # Create retriever if not provided + if retriever is None: + self.retriever = CLIPRetriever() + else: + self.retriever = retriever + + def build_few_shot_prompt( + self, + base_prompt: str, + retrieved_metadata: List[Dict], + format_fn=None + ) -> str: + """ + Build few-shot prompt from retrieved examples + + Args: + base_prompt: Original user prompt + retrieved_metadata: List of metadata dicts from retrieval + format_fn: Optional function to format each example + + Returns: + Augmented prompt with examples + """ + if format_fn is None: + # Default formatting + format_fn = lambda meta: f"Example: {meta.get('description', 'No description')}\n" + + # Build context from retrieved examples + context = "Here are some similar examples:\n\n" + for i, meta in enumerate(retrieved_metadata, 1): + context += f"{i}. {format_fn(meta)}\n" + + # Combine with user prompt + augmented_prompt = f"{context}\nNow, {base_prompt}" + + return augmented_prompt + + def generate( + self, + image: Union[str, Image.Image], + prompt: str, + use_rag: bool = True, + max_new_tokens: int = 256, + temperature: float = 0.7, + top_k: int = 50, + top_p: float = 0.9, + format_fn=None + ) -> Dict: + """ + Generate response with optional RAG + + Args: + image: Query image or path + prompt: User prompt + use_rag: Whether to use retrieval-augmented generation + max_new_tokens: Max tokens to generate + temperature: Sampling temperature + top_k: Top-k sampling + top_p: Nucleus sampling + format_fn: Optional function to format retrieved examples + + Returns: + Dict with 'response', 'retrieved_metadata', 'augmented_prompt' + """ + # Load image if path + if isinstance(image, str): + image = Image.open(image) + + retrieved_metadata = [] + augmented_prompt = prompt + + # Retrieve if RAG enabled + if use_rag: + distances, retrieved_metadata = self.retriever.retrieve( + image, + self.vector_store, + k=self.k + ) + + # Build few-shot prompt + augmented_prompt = self.build_few_shot_prompt( + prompt, + retrieved_metadata, + format_fn=format_fn + ) + + print(f"\nšŸ“š Retrieved {len(retrieved_metadata)} examples") + print(f" Distances: {distances}") + + # Generate with VLM + # Note: Actual generation code depends on MiniMindVLM's interface + # This is a placeholder showing the concept + + # For now, return structure (actual generation would call model.generate()) + return { + 'response': augmented_prompt, # Placeholder: would be actual generation + 'retrieved_metadata': retrieved_metadata, + 'augmented_prompt': augmented_prompt, + 'retrieval_distances': distances.tolist() if use_rag else [] + } + + def evaluate_retrieval( + self, + test_images: List[Union[str, Image.Image]], + ground_truth_labels: List[str], + metadata_label_fn=lambda meta: meta.get('label', '') + ) -> Dict: + """ + Evaluate retrieval quality + + Args: + test_images: List of test images + ground_truth_labels: True labels for each test image + metadata_label_fn: Function to extract label from metadata + + Returns: + Dict with Recall@K, MRR, etc. + """ + recall_at_k = 0 + mrr = 0 + + for image, true_label in zip(test_images, ground_truth_labels): + distances, retrieved_metadata = self.retriever.retrieve( + image, + self.vector_store, + k=self.k + ) + + # Check if true label in retrieved results + retrieved_labels = [metadata_label_fn(meta) for meta in retrieved_metadata] + + if true_label in retrieved_labels: + recall_at_k += 1 + # Mean reciprocal rank + rank = retrieved_labels.index(true_label) + 1 + mrr += 1.0 / rank + + recall_at_k /= len(test_images) + mrr /= len(test_images) + + return { + f'Recall@{self.k}': recall_at_k, + 'MRR': mrr + } + + +# Example usage +if __name__ == "__main__": + from transformers import AutoTokenizer + + # Load VLM + from model.model_vlm import VLMConfig + config = VLMConfig() + vlm = MiniMindVLM(config) + tokenizer = AutoTokenizer.from_pretrained('./model') + + # Load vector store (assuming it exists) + # vector_store = VectorStore.load("knowledge_base.faiss") + + # Create RAG-VLM + # rag_vlm = RAG_VLM(vlm, tokenizer, vector_store, k=3) + + # Generate with RAG + # result = rag_vlm.generate( + # image="test_image.jpg", + # prompt="Describe this scene in detail", + # use_rag=True + # ) + + print("RAG-VLM module ready!") diff --git a/rag/retriever.py b/rag/retriever.py new file mode 100644 index 0000000..b1733b2 --- /dev/null +++ b/rag/retriever.py @@ -0,0 +1,185 @@ +""" +CLIP-based Image Retriever for RAG System +Uses frozen CLIP ViT-B/16 from MiniMindVLM for visual similarity +""" +import os +import sys +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import torch +import numpy as np +from PIL import Image +from typing import List, Union, Dict, Tuple +from model.model_vlm import MiniMindVLM +from rag.vector_store import VectorStore + + +class CLIPRetriever: + """ + CLIP-based retriever for visual similarity search + + Uses the same CLIP ViT-B/16 model as MiniMindVLM for consistency + """ + + def __init__( + self, + vision_model_path: str = "./model/vision_model/clip-vit-base-patch16", + device: str = "cuda" if torch.cuda.is_available() else "cpu" + ): + """ + Initialize CLIP retriever + + Args: + vision_model_path: Path to CLIP model + device: Device to run on + """ + self.device = device + + # Load CLIP model (reuse MiniMindVLM's CLIP loader) + self.vision_model, self.processor = MiniMindVLM.get_vision_model(vision_model_path) + + if self.vision_model is None: + raise ValueError(f"CLIP model not found at {vision_model_path}") + + self.vision_model = self.vision_model.to(device) + self.vision_model.eval() + + print(f"āœ… Loaded CLIP ViT-B/16 on {device}") + + def encode_image(self, image: Union[str, Image.Image]) -> np.ndarray: + """ + Encode a single image to CLIP embedding + + Args: + image: PIL Image or path to image + + Returns: + Embedding as numpy array of shape (768,) + """ + # Load image if path provided + if isinstance(image, str): + image = Image.open(image) + + # Convert image to tensor + image_tensor = MiniMindVLM.image2tensor(image, self.processor).to(self.device) + + # Get CLIP embedding + with torch.no_grad(): + embedding = MiniMindVLM.get_image_embeddings(image_tensor, self.vision_model) + # Pool patch embeddings to single vector + embedding = embedding.mean(dim=0).cpu().numpy() # Shape: (768,) + + return embedding + + def encode_images_batch(self, images: List[Union[str, Image.Image]], batch_size: int = 32) -> np.ndarray: + """ + Encode multiple images in batches + + Args: + images: List of PIL Images or paths + batch_size: Batch size for encoding + + Returns: + Embeddings as numpy array of shape (N, 768) + """ + embeddings = [] + + for i in range(0, len(images), batch_size): + batch = images[i:i + batch_size] + + # Load images + batch_images = [] + for img in batch: + if isinstance(img, str): + img = Image.open(img) + batch_images.append(img) + + # Convert to tensors + # image2tensor returns (1, C, H, W), so we concat on batch dim + batch_tensors = torch.cat([ + MiniMindVLM.image2tensor(img, self.processor) + for img in batch_images + ], dim=0).to(self.device) + + # Encode + with torch.no_grad(): + batch_embeddings = MiniMindVLM.get_image_embeddings(batch_tensors, self.vision_model) + # Pool: (batch, patches, 768) -> (batch, 768) + batch_embeddings = batch_embeddings.mean(dim=1).cpu().numpy() + + embeddings.append(batch_embeddings) + + return np.vstack(embeddings) + + def build_index( + self, + images: List[Union[str, Image.Image]], + metadata: List[Dict], + save_path: str, + batch_size: int = 32 + ) -> VectorStore: + """ + Build FAISS index from images + + Args: + images: List of images or paths + metadata: List of metadata dicts (one per image) + save_path: Where to save index + batch_size: Encoding batch size + + Returns: + VectorStore with built index + """ + print(f"šŸ”Ø Building index from {len(images)} images...") + + # Encode all images + embeddings = self.encode_images_batch(images, batch_size=batch_size) + print(f"āœ… Encoded {len(embeddings)} images to {embeddings.shape[1]}-dim vectors") + + # Build vector store + store = VectorStore(dimension=embeddings.shape[1]) + store.add(embeddings, metadata) + + # Save + store.save(save_path) + + return store + + def retrieve( + self, + query_image: Union[str, Image.Image], + vector_store: VectorStore, + k: int = 5 + ) -> Tuple[np.ndarray, List[Dict]]: + """ + Retrieve k most similar images from index + + Args: + query_image: Query image or path + vector_store: VectorStore to search in + k: Number of results to return + + Returns: + distances: Array of shape (k,) with distances + metadata: List of k metadata dicts + """ + # Encode query + query_embedding = self.encode_image(query_image) + query_embedding = query_embedding.reshape(1, -1) # (1, 768) + + # Search + distances, indices, metadata = vector_store.search(query_embedding, k=k) + + # Return flattened results (since we only have 1 query) + return distances[0], metadata[0] + + +# Example usage +if __name__ == "__main__": + # Create retriever + retriever = CLIPRetriever() + + # Example: encode single image + # img_path = "dataset/eval_images/åŸŽåø‚č½¦ę°“é©¬é¾™-city-traffic.jpg" + # embedding = retriever.encode_image(img_path) + # print(f"Image embedding shape: {embedding.shape}") diff --git a/rag/vector_store.py b/rag/vector_store.py new file mode 100644 index 0000000..f2bd318 --- /dev/null +++ b/rag/vector_store.py @@ -0,0 +1,193 @@ +""" +FAISS-based Vector Store for Image Embeddings +Enables fast similarity search over large image collections +""" +import os +import json +import numpy as np +import faiss +from typing import List, Dict, Optional, Tuple + + +class VectorStore: + """ + FAISS-based vector store for efficient similarity search + + Supports: + - L2 (Euclidean) distance search + - Incremental index building + - Save/load to disk + - Metadata storage alongside embeddings + """ + + def __init__(self, dimension: int = 768, metric: str = "L2"): + """ + Initialize vector store + + Args: + dimension: Embedding dimension (768 for CLIP ViT-B/16) + metric: Distance metric ("L2" or "IP" for inner product) + """ + self.dimension = dimension + self.metric = metric + + # Initialize FAISS index + if metric == "L2": + self.index = faiss.IndexFlatL2(dimension) + elif metric == "IP": + self.index = faiss.IndexFlatIP(dimension) # Inner product (cosine if normalized) + else: + raise ValueError(f"Unknown metric: {metric}. Use 'L2' or 'IP'") + + # Metadata storage (parallel to FAISS index) + self.metadata = [] + + def add(self, embeddings: np.ndarray, metadata: Optional[List[Dict]] = None): + """ + Add embeddings to the index + + Args: + embeddings: Array of shape (N, dimension) + metadata: List of N metadata dicts (one per embedding) + """ + if embeddings.dtype != np.float32: + embeddings = embeddings.astype(np.float32) + + # Normalize for inner product (makes it cosine similarity) + if self.metric == "IP": + faiss.normalize_L2(embeddings) + + # Add to FAISS index + self.index.add(embeddings) + + # Add metadata + if metadata: + assert len(metadata) == len(embeddings), "Metadata length must match embeddings" + self.metadata.extend(metadata) + else: + # Add empty dicts if no metadata provided + self.metadata.extend([{} for _ in range(len(embeddings))]) + + def search(self, query_embeddings: np.ndarray, k: int = 5) -> Tuple[np.ndarray, np.ndarray, List[List[Dict]]]: + """ + Search for k nearest neighbors + + Args: + query_embeddings: Array of shape (N_queries, dimension) + k: Number of neighbors to return + + Returns: + distances: Array of shape (N_queries, k) + indices: Array of shape (N_queries, k) + metadata: List of lists, metadata for each query's k neighbors + """ + if query_embeddings.dtype != np.float32: + query_embeddings = query_embeddings.astype(np.float32) + + # Normalize for inner product + if self.metric == "IP": + faiss.normalize_L2(query_embeddings) + + # Search + distances, indices = self.index.search(query_embeddings, k) + + # Retrieve metadata + retrieved_metadata = [] + for query_indices in indices: + query_metadata = [self.metadata[i] if i < len(self.metadata) else {} for i in query_indices] + retrieved_metadata.append(query_metadata) + + return distances, indices, retrieved_metadata + + def save(self, index_path: str, metadata_path: Optional[str] = None): + """ + Save index and metadata to disk + + Args: + index_path: Path to save FAISS index (.faiss or .index) + metadata_path: Path to save metadata (.json). If None, uses index_path.json + """ + # Save FAISS index + faiss.write_index(self.index, index_path) + + # Save metadata + if metadata_path is None: + metadata_path = index_path.replace('.faiss', '.json').replace('.index', '.json') + if not metadata_path.endswith('.json'): + metadata_path += '.json' + + with open(metadata_path, 'w') as f: + json.dump({ + 'metadata': self.metadata, + 'dimension': self.dimension, + 'metric': self.metric, + 'count': len(self.metadata) + }, f, indent=2) + + print(f"āœ… Saved FAISS index to {index_path}") + print(f"āœ… Saved metadata to {metadata_path}") + + @classmethod + def load(cls, index_path: str, metadata_path: Optional[str] = None) -> 'VectorStore': + """ + Load index and metadata from disk + + Args: + index_path: Path to FAISS index + metadata_path: Path to metadata. If None, uses index_path.json + + Returns: + VectorStore instance with loaded data + """ + # Load FAISS index + index = faiss.read_index(index_path) + + # Load metadata + if metadata_path is None: + metadata_path = index_path.replace('.faiss', '.json').replace('.index', '.json') + if not metadata_path.endswith('.json'): + metadata_path += '.json' + + with open(metadata_path, 'r') as f: + data = json.load(f) + + # Create instance + store = cls(dimension=data['dimension'], metric=data['metric']) + store.index = index + store.metadata = data['metadata'] + + print(f"āœ… Loaded {data['count']} embeddings from {index_path}") + + return store + + def __len__(self): + """Return number of vectors in index""" + return self.index.ntotal + + +# Example usage +if __name__ == "__main__": + # Create some random embeddings for demo + dimension = 768 + n_vectors = 100 + embeddings = np.random.randn(n_vectors, dimension).astype(np.float32) + + # Create metadata + metadata = [{"id": f"img_{i}", "description": f"Image {i}"} for i in range(n_vectors)] + + # Build index + store = VectorStore(dimension=dimension) + store.add(embeddings, metadata) + + # Search + query = np.random.randn(1, dimension).astype(np.float32) + distances, indices, retrieved_metadata = store.search(query, k=5) + + print(f"Top 5 matches:") + for i, (dist, idx, meta) in enumerate(zip(distances[0], indices[0], retrieved_metadata[0])): + print(f" {i+1}. Index {idx}, Distance {dist:.4f}, Metadata: {meta}") + + # Save and load + store.save("demo_index.faiss") + loaded_store = VectorStore.load("demo_index.faiss") + print(f"Loaded store has {len(loaded_store)} vectors") diff --git a/rag_demo_index.faiss b/rag_demo_index.faiss new file mode 100644 index 0000000..693d2fd Binary files /dev/null and b/rag_demo_index.faiss differ diff --git a/rag_demo_index.json b/rag_demo_index.json new file mode 100644 index 0000000..5a4e514 --- /dev/null +++ b/rag_demo_index.json @@ -0,0 +1,55 @@ +{ + "metadata": [ + { + "id": "\u821e\u8e48-dance", + "filename": "\u821e\u8e48-dance.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u821e\u8e48-dance.jpg", + "description": "Image: \u821e\u8e48-dance.jpg" + }, + { + "id": "\u81ea\u884c\u8f66\u9c9c\u82b1-Bicycle-Flowers", + "filename": "\u81ea\u884c\u8f66\u9c9c\u82b1-Bicycle-Flowers.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u81ea\u884c\u8f66\u9c9c\u82b1-Bicycle-Flowers.jpg", + "description": "Image: \u81ea\u884c\u8f66\u9c9c\u82b1-Bicycle-Flowers.jpg" + }, + { + "id": "\u718a\u732b\u8349\u5730-Panda-Grassland", + "filename": "\u718a\u732b\u8349\u5730-Panda-Grassland.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u718a\u732b\u8349\u5730-Panda-Grassland.jpg", + "description": "Image: \u718a\u732b\u8349\u5730-Panda-Grassland.jpg" + }, + { + "id": "\u6905\u5b50\u8001\u4eba\u770b\u4e66-Chair-Elderly-Reading", + "filename": "\u6905\u5b50\u8001\u4eba\u770b\u4e66-Chair-Elderly-Reading.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u6905\u5b50\u8001\u4eba\u770b\u4e66-Chair-Elderly-Reading.jpg", + "description": "Image: \u6905\u5b50\u8001\u4eba\u770b\u4e66-Chair-Elderly-Reading.jpg" + }, + { + "id": "\u5f69\u8679\u7011\u5e03-Rainbow-Falls", + "filename": "\u5f69\u8679\u7011\u5e03-Rainbow-Falls.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u5f69\u8679\u7011\u5e03-Rainbow-Falls.jpg", + "description": "Image: \u5f69\u8679\u7011\u5e03-Rainbow-Falls.jpg" + }, + { + "id": "\u5c0f\u72d7\u7f8e\u5973\u6d77\u8fb9-Dog-Woman-Sea", + "filename": "\u5c0f\u72d7\u7f8e\u5973\u6d77\u8fb9-Dog-Woman-Sea.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u5c0f\u72d7\u7f8e\u5973\u6d77\u8fb9-Dog-Woman-Sea.jpg", + "description": "Image: \u5c0f\u72d7\u7f8e\u5973\u6d77\u8fb9-Dog-Woman-Sea.jpg" + }, + { + "id": "\u592a\u7a7a\u5b87\u822a\u5458-Astronaut-Space", + "filename": "\u592a\u7a7a\u5b87\u822a\u5458-Astronaut-Space.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u592a\u7a7a\u5b87\u822a\u5458-Astronaut-Space.jpg", + "description": "Image: \u592a\u7a7a\u5b87\u822a\u5458-Astronaut-Space.jpg" + }, + { + "id": "\u57ce\u5e02\u8f66\u6c34\u9a6c\u9f99-city-traffic", + "filename": "\u57ce\u5e02\u8f66\u6c34\u9a6c\u9f99-city-traffic.jpg", + "path": "/workspace/minimind-v/minimind-v/dataset/eval_images/\u57ce\u5e02\u8f66\u6c34\u9a6c\u9f99-city-traffic.jpg", + "description": "Image: \u57ce\u5e02\u8f66\u6c34\u9a6c\u9f99-city-traffic.jpg" + } + ], + "dimension": 768, + "metric": "L2", + "count": 8 +} \ No newline at end of file diff --git a/scripts/build_knowledge_base.py b/scripts/build_knowledge_base.py new file mode 100755 index 0000000..d3af62e --- /dev/null +++ b/scripts/build_knowledge_base.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +""" +Build RAG Knowledge Base from Images +Creates FAISS index from image directory for retrieval-augmented generation +""" +import os +import sys +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import argparse +import json +from pathlib import Path +from tqdm import tqdm +from PIL import Image +from rag.retriever import CLIPRetriever + + +def build_knowledge_base( + images_dir: str, + output_path: str, + description_file: str = None, + batch_size: int = 32 +): + """ + Build knowledge base from image directory + + Args: + images_dir: Directory containing images + output_path: Where to save FAISS index + description_file: Optional JSON file with descriptions per image + batch_size: Batch size for encoding + """ + print(f"\n{'='*60}") + print("šŸ”Ø Building Knowledge Base for RAG") + print(f"{'='*60}\n") + + # Find all images + image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'} + images_dir = Path(images_dir) + image_files = [ + f for f in images_dir.iterdir() + if f.suffix.lower() in image_extensions + ] + + print(f"šŸ“‚ Found {len(image_files)} images in {images_dir}") + + # Load descriptions if provided + descriptions = {} + if description_file and os.path.exists(description_file): + with open(description_file, 'r') as f: + descriptions = json.load(f) + print(f"šŸ“ Loaded descriptions for {len(descriptions)} images") + + # Create metadata + metadata = [] + image_paths = [] + + for img_file in tqdm(image_files, desc="Preparing metadata"): + img_name = img_file.name + + # Create metadata entry + meta = { + 'id': img_file.stem, + 'filename': img_name, + 'path': str(img_file.absolute()), + 'description': descriptions.get(img_name, f"Image: {img_name}"), + } + + metadata.append(meta) + image_paths.append(str(img_file)) + + # Create retriever + print("\nšŸ”§ Initializing CLIP retriever...") + retriever = CLIPRetriever() + + # Build index + print(f"\nšŸš€ Encoding images in batches of {batch_size}...") + vector_store = retriever.build_index( + images=image_paths, + metadata=metadata, + save_path=output_path, + batch_size=batch_size + ) + + print(f"\n{'='*60}") + print(f"āœ… Knowledge Base Built Successfully!") + print(f"{'='*60}") + print(f" Total images: {len(image_paths)}") + print(f" Index file: {output_path}") + print(f" Metadata file: {output_path.replace('.faiss', '.json')}") + print(f" Embedding dimension: {vector_store.dimension}") + print(f"{'='*60}\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Build RAG knowledge base from images") + parser.add_argument( + "--images_dir", + type=str, + required=True, + help="Directory containing images" + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="Output path for FAISS index (.faiss)" + ) + parser.add_argument( + "--descriptions", + type=str, + default=None, + help="Optional JSON file with image descriptions" + ) + parser.add_argument( + "--batch_size", + type=int, + default=32, + help="Batch size for encoding" + ) + + args = parser.parse_args() + + build_knowledge_base( + images_dir=args.images_dir, + output_path=args.output, + description_file=args.descriptions, + batch_size=args.batch_size + ) diff --git a/scripts/demo_rag.py b/scripts/demo_rag.py new file mode 100755 index 0000000..8700e09 --- /dev/null +++ b/scripts/demo_rag.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +""" +RAG System Demo +Demonstrates retrieval-augmented generation with MiniMindVLM +""" +import os +import sys +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import argparse +from PIL import Image +from rag.retriever import CLIPRetriever +from rag.vector_store import VectorStore + + +def demo_retrieval(query_image_path: str, index_path: str, k: int = 3): + """ + Demo RAG retrieval + + Args: + query_image_path: Path to query image + index_path: Path to FAISS index + k: Number of results to retrieve + """ + print(f"\n{'='*70}") + print("šŸ” RAG Retrieval Demo") + print(f"{'='*70}\n") + + # Load query image + query_image = Image.open(query_image_path) + print(f"šŸ“· Query image: {query_image_path}") + print(f" Size: {query_image.size}") + print() + + # Load vector store + print(f"šŸ“‚ Loading knowledge base from {index_path}...") + vector_store = VectorStore.load(index_path) + print() + + # Create retriever + print("šŸ”§ Initializing CLIP retriever...") + retriever = CLIPRetriever() + print() + + # Retrieve similar images + print(f"šŸ” Retrieving top-{k} similar images...") + distances, metadata = retriever.retrieve(query_image, vector_store, k=k) + print() + + # Display results + print(f"{'='*70}") + print(f"šŸ“Š Top-{k} Most Similar Images:") + print(f"{'='*70}\n") + + for i, (dist, meta) in enumerate(zip(distances, metadata), 1): + print(f"{i}. {meta['filename']}") + print(f" Distance: {dist:.4f}") + print(f" Description: {meta.get('description', 'No description')}") + print(f" Path: {meta.get('path', 'Unknown')}") + print() + + print(f"{'='*70}\n") + + # Recommendation + if distances[0] < 5.0: + print("āœ… Retrieved very similar images (distance < 5.0)") + elif distances[0] < 10.0: + print("āš ļø Retrieved moderately similar images (distance < 10.0)") + else: + print("āŒ No very similar images found (distance >= 10.0)") + + print(f"\n{'='*70}\n") + + return distances, metadata + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Demo RAG retrieval") + parser.add_argument( + "--query", + type=str, + required=True, + help="Path to query image" + ) + parser.add_argument( + "--index", + type=str, + default="rag_demo_index.faiss", + help="Path to FAISS index" + ) + parser.add_argument( + "--k", + type=int, + default=3, + help="Number of results to retrieve" + ) + + args = parser.parse_args() + + demo_retrieval(args.query, args.index, args.k) diff --git a/trainer/lora_utils.py b/trainer/lora_utils.py new file mode 100644 index 0000000..734c633 --- /dev/null +++ b/trainer/lora_utils.py @@ -0,0 +1,144 @@ +""" +LoRA (Low-Rank Adaptation) utilities for MiniMindVLM +Enables parameter-efficient fine-tuning with <2% trainable parameters +""" +import torch +from peft import LoraConfig, get_peft_model, TaskType, PeftModel +from typing import Optional + + +def get_lora_config( + r: int = 16, + lora_alpha: int = 32, + lora_dropout: float = 0.05, + target_modules: Optional[list] = None, + bias: str = "none", +) -> LoraConfig: + """ + Get LoRA configuration for MiniMindVLM + + Args: + r: Rank of LoRA update matrices (higher = more parameters) + lora_alpha: Scaling factor for LoRA updates + lora_dropout: Dropout rate for LoRA layers + target_modules: Which modules to apply LoRA to. If None, uses default for MiniMind + bias: How to handle bias parameters ("none", "all", "lora_only") + + Returns: + LoraConfig object ready to apply to model + + Typical memory savings: + - r=8: ~0.5% trainable params, 99.5% frozen + - r=16: ~1.0% trainable params, 99.0% frozen (recommended) + - r=32: ~2.0% trainable params, 98.0% frozen + """ + + # Default target modules for MiniMind architecture + # These are the attention projection layers + if target_modules is None: + target_modules = [ + "c_attn", # Combined QKV projection in attention + "c_proj", # Output projection in attention + "c_fc", # First linear in MLP (optional, but helps) + ] + + config = LoraConfig( + task_type=TaskType.CAUSAL_LM, + inference_mode=False, + r=r, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + target_modules=target_modules, + bias=bias, + ) + + return config + + +def apply_lora(model, lora_config: LoraConfig, verbose: bool = True): + """ + Apply LoRA adapters to a MiniMindVLM model + + Args: + model: MiniMindVLM model instance + lora_config: LoRA configuration from get_lora_config() + verbose: Print trainable parameter statistics + + Returns: + Model with LoRA adapters applied + """ + # Freeze vision encoder (it's already frozen in MiniMindVLM, but be explicit) + if hasattr(model, 'vision_encoder') and model.vision_encoder is not None: + for param in model.vision_encoder.parameters(): + param.requires_grad = False + + # Apply LoRA + model = get_peft_model(model, lora_config) + + if verbose: + model.print_trainable_parameters() + + # Calculate memory savings + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + total = sum(p.numel() for p in model.parameters()) + + print(f"\nšŸ“Š LoRA Statistics:") + print(f" Total parameters: {total / 1e6:.2f}M") + print(f" Trainable parameters: {trainable / 1e6:.2f}M ({trainable / total * 100:.2f}%)") + print(f" Frozen parameters: {(total - trainable) / 1e6:.2f}M ({(total - trainable) / total * 100:.2f}%)") + print(f" Memory savings: ~{100 - (trainable / total * 100):.1f}%") + + return model + + +def load_lora_model(base_model, lora_adapter_path: str): + """ + Load a model with LoRA adapters + + Args: + base_model: Base MiniMindVLM model (without LoRA) + lora_adapter_path: Path to saved LoRA adapters + + Returns: + Model with LoRA adapters loaded + """ + model = PeftModel.from_pretrained(base_model, lora_adapter_path) + return model + + +def merge_lora_weights(model, save_path: Optional[str] = None): + """ + Merge LoRA weights into base model for deployment + After merging, the model is a standard model without adapters + + Args: + model: Model with LoRA adapters + save_path: If provided, save merged model to this path + + Returns: + Merged model (standard model, no adapters) + """ + # Merge LoRA weights into base model + model = model.merge_and_unload() + + if save_path: + model.save_pretrained(save_path) + print(f"āœ… Merged model saved to {save_path}") + + return model + + +# Example usage: +if __name__ == "__main__": + from model.model_vlm import MiniMindVLM, VLMConfig + + # Create model + config = VLMConfig() + model = MiniMindVLM(config) + + # Apply LoRA + lora_config = get_lora_config(r=16) + model = apply_lora(model, lora_config) + + # Now model is ready for training with LoRA + # Only ~1% of parameters will be updated during training diff --git a/trainer/train_sft_vlm_lora.py b/trainer/train_sft_vlm_lora.py new file mode 100644 index 0000000..57707e7 --- /dev/null +++ b/trainer/train_sft_vlm_lora.py @@ -0,0 +1,232 @@ +""" +SFT VLM Training with LoRA Support +Extension of train_sft_vlm.py with LoRA for parameter-efficient fine-tuning +""" +import os +import sys + +__package__ = "trainer" +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import argparse +import time +import warnings +import torch +import torch.distributed as dist +from contextlib import nullcontext +from torch import optim, nn +from torch.nn.parallel import DistributedDataParallel +from torch.utils.data import DataLoader, DistributedSampler +from transformers import AutoTokenizer +from model.model_vlm import MiniMindVLM, VLMConfig +from dataset.lm_dataset import VLMDataset +from trainer.trainer_utils import get_lr, Logger, is_main_process, init_distributed_mode, setup_seed, init_vlm_model, vlm_checkpoint, SkipBatchSampler +from trainer.lora_utils import get_lora_config, apply_lora + +warnings.filterwarnings('ignore') + + +def train_epoch(epoch, loader, iters, start_step=0, wandb=None): + start_time = time.time() + for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1): + input_ids = input_ids.to(args.device) + labels = labels.to(args.device) + pixel_values = pixel_values.to(args.device) + lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) + for param_group in optimizer.param_groups: + param_group['lr'] = lr + + with autocast_ctx: + res = model(input_ids, labels=labels, pixel_values=pixel_values) + loss = res.loss + res.aux_loss + loss = loss / args.accumulation_steps + + scaler.scale(loss).backward() + + if (step + 1) % args.accumulation_steps == 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + + scaler.step(optimizer) + scaler.update() + + optimizer.zero_grad(set_to_none=True) + + if step % args.log_interval == 0 or step == iters - 1: + spend_time = time.time() - start_time + logger.log( + 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.8f} epoch_time:{}min:'.format( + epoch + 1, + args.epochs, + step, + iters, + loss.item() * args.accumulation_steps, + optimizer.param_groups[-1]['lr'], + spend_time / (step + 1) * iters // 60 - spend_time // 60)) + + if (wandb is not None) and is_main_process(): + wandb.log({"loss": loss.item() * args.accumulation_steps, + "lr": optimizer.param_groups[-1]['lr'], + "epoch": epoch + 1, + "step": step}) + + if step % args.save_interval == 0 or step == iters - 1: + if is_main_process(): + vlm_checkpoint(model.module if hasattr(model, 'module') else model, + epoch, + step, + args, + suffix=f"_lora" if args.use_lora else "") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="MiniMind VLM SFT with LoRA") + parser.add_argument("--out_dir", type=str, default="out", help="Output directory") + parser.add_argument("--epochs", type=int, default=1, help="Number of epochs") + parser.add_argument("--batch_size", type=int, default=64, help="Batch size") + parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") + parser.add_argument("--dtype", type=str, default="bfloat16" if torch.cuda.is_bf16_supported() else "float16") + parser.add_argument("--use_wandb", action="store_true", help="Use Weights & Biases") + parser.add_argument("--num_workers", type=int, default=0, help="Data loader workers") + parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation") + parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping") + parser.add_argument("--warmup_iters", type=int, default=0, help="Warmup iterations") + parser.add_argument("--log_interval", type=int, default=100, help="Log interval") + parser.add_argument("--save_interval", type=int, default=1000, help="Save interval") + parser.add_argument("--local_rank", type=int, default=-1, help="DDP local rank") + parser.add_argument("--use_moe", action="store_true", help="Use Mixture of Experts") + parser.add_argument("--resume_step", type=int, default=0, help="Resume from step") + parser.add_argument("--hidden_size", type=int, default=512, help="Hidden size") + parser.add_argument("--num_heads", type=int, default=16, help="Number of attention heads") + parser.add_argument("--num_layers", type=int, default=8, help="Number of layers") + + # LoRA-specific arguments + parser.add_argument("--use_lora", action="store_true", help="Use LoRA for parameter-efficient training") + parser.add_argument("--lora_r", type=int, default=16, help="LoRA rank") + parser.add_argument("--lora_alpha", type=int, default=32, help="LoRA alpha") + parser.add_argument("--lora_dropout", type=float, default=0.05, help="LoRA dropout") + + args = parser.parse_args() + + # Setup + init_distributed_mode(args) + setup_seed(42) + logger = Logger(args.out_dir, is_main_process()) + + # Load tokenizer and model + tokenizer = AutoTokenizer.from_pretrained('./model') + + # Create VLM config + lm_config = VLMConfig( + hidden_size=args.hidden_size, + num_attention_heads=args.num_heads, + num_hidden_layers=args.num_layers, + use_moe=args.use_moe, + ) + + # Initialize model + model = init_vlm_model(lm_config, args.device) + + # Apply LoRA if requested + if args.use_lora: + print("\n" + "=" * 50) + print("šŸš€ Applying LoRA for Parameter-Efficient Training") + print("=" * 50) + + lora_config = get_lora_config( + r=args.lora_r, + lora_alpha=args.lora_alpha, + lora_dropout=args.lora_dropout, + ) + + model = apply_lora(model, lora_config, verbose=True) + print("=" * 50 + "\n") + + # DDP if needed + if args.ddp: + model = DistributedDataParallel(model, device_ids=[args.local_rank]) + + # Optimizer + optimizer = optim.Adam( + filter(lambda p: p.requires_grad, model.parameters()), # Only train LoRA params if using LoRA + lr=args.learning_rate + ) + + # AMP scaler + scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16'])) + autocast_ctx = ( + torch.cuda.amp.autocast(dtype=torch.bfloat16) + if args.dtype == 'bfloat16' + else torch.cuda.amp.autocast(dtype=torch.float16) + if args.dtype == 'float16' + else nullcontext() + ) + + # Load dataset + dataset = VLMDataset( + parquet_path='./dataset/sft_vlm.parquet', + tokenizer=tokenizer, + preprocess=model.processor if not args.ddp else model.module.processor, + max_length=512 + ) + + # Data loader + if args.ddp: + sampler = DistributedSampler(dataset) + loader = DataLoader( + dataset, + batch_size=args.batch_size, + sampler=sampler, + num_workers=args.num_workers, + pin_memory=True + ) + else: + loader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=True + ) + + # WandB + wandb = None + if args.use_wandb and is_main_process(): + import wandb as wandb_lib + wandb = wandb_lib + wandb.init( + project="drivemind-v", + name=f"sft_lora_r{args.lora_r}" if args.use_lora else "sft_full", + config=vars(args) + ) + + # Training loop + iters = len(loader) + logger.log(f"Starting SFT training {'with LoRA' if args.use_lora else 'full fine-tuning'}...") + logger.log(f"Total iterations: {iters} per epoch") + + for epoch in range(args.epochs): + train_epoch(epoch, loader, iters, start_step=0, wandb=wandb) + + # Save final model + if is_main_process(): + if args.use_lora: + # Save LoRA adapters (tiny file ~2-5MB) + save_path = f"{args.out_dir}/lora_adapters_sft_vlm_{args.hidden_size}" + model_to_save = model.module if hasattr(model, 'module') else model + model_to_save.save_pretrained(save_path) + logger.log(f"āœ… LoRA adapters saved to {save_path}") + else: + vlm_checkpoint( + model.module if hasattr(model, 'module') else model, + args.epochs - 1, + iters - 1, + args, + suffix="_final" + ) + + if wandb: + wandb.finish() + + logger.log("Training complete!")