ADD: pipeline

data_processing/data_loader.py

@@ -1,74 +1,29 @@
-from typing import Dict, List, Optional
+
 import logging
+from datasets import load_dataset
 
 logger = logging.getLogger(__name__)
 
 class DataLoader:
     def __init__(self, cache_dir: str = "./cache"):
         self.cache_dir = cache_dir
-        
-    def load_hotpotqa(self, split: str = "train"):
-        """Load HotpotQA dataset for multi-hop reasoning (simplified version)"""
-        try:
-            # Simplified version - return empty list for demo
-            logger.info(f"Loading HotpotQA {split} (simplified version)")
-            return []
-        except Exception as e:
-            logger.error(f"Failed to load HotpotQA: {e}")
-            raise
-    
-    def load_triviaqa(self, split: str = "train"):
-        """Load TriviaQA dataset for open-domain QA (simplified version)"""
-        try:
-            logger.info(f"Loading TriviaQA {split} (simplified version)")
-            return []
-        except Exception as e:
-            logger.error(f"Failed to load TriviaQA: {e}")
-            raise
-    
-    def load_wikipedia(self, language: str = "en", date: str = "20231101"):
-        """Load Wikipedia dump for knowledge base (simplified version)"""
-        try:
-            logger.info(f"Loading Wikipedia {language} (simplified version)")
-            return []
-        except Exception as e:
-            logger.error(f"Failed to load Wikipedia: {e}")
-            raise
-    
-    def load_nq_open(self, split: str = "train"):
-        """Load Natural Questions Open dataset (simplified version)"""
-        try:
-            logger.info(f"Loading NQ Open {split} (simplified version)")
-            return []
-        except Exception as e:
-            logger.error(f"Failed to load NQ Open: {e}")
-            raise
-    
-    def get_qa_datasets(self) -> Dict[str, List]:
-        """Load all QA datasets (simplified version)"""
-        datasets = {}
+
+    def load_msmarco_passage(self, split: str = "train"):
+        """Load MS MARCO Passage Ranking dataset from Hugging Face (v2.1)"""
         try:
-            datasets['hotpotqa'] = self.load_hotpotqa()
-            datasets['triviaqa'] = self.load_triviaqa()
-            datasets['nq_open'] = self.load_nq_open()
-            logger.info("All QA datasets loaded successfully")
-            return datasets
+            logger.info(f"Downloading MS MARCO Passage Ranking {split} (v2.1) from Hugging Face")
+            ds = load_dataset("ms_marco", "v2.1", split=split)
+            return ds
         except Exception as e:
-            logger.error(f"Failed to load QA datasets: {e}")
+            logger.error(f"Failed to load MS MARCO Passage Ranking: {e}")
             raise
-    
-    def get_knowledge_base(self) -> List[str]:
-        """Load knowledge base (simplified version)"""
+
+    def get_passage_dataset(self, split: str = "train"):
+        """Load MS MARCO Passage Ranking dataset"""
         try:
-            logger.info("Loading knowledge base (simplified version)")
-            # Return some sample passages for demo
-            return [
-                "Machine learning is a subset of artificial intelligence that focuses on algorithms.",
-                "The capital of France is Paris.",
-                "Python is a popular programming language used for data science.",
-                "The Great Wall of China is one of the most famous landmarks in the world.",
-                "Climate change refers to long-term shifts in global temperatures and weather patterns."
-            ]
+            ds = self.load_msmarco_passage(split)
+            logger.info("MS MARCO Passage Ranking loaded successfully")
+            return ds
         except Exception as e:
-            logger.error(f"Failed to load knowledge base: {e}")
+            logger.error(f"Failed to load MS MARCO Passage Ranking: {e}")
             raise

data_processing/preprocessor.py

@@ -66,17 +66,25 @@ class Preprocessor:
         return processed
     
     def preprocess_qa_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        """Preprocess QA data"""
+        """Preprocess QA data, auto convert dict/list fields to string"""
         processed = []
-        
+        def to_str(val):
+            if isinstance(val, dict):
+                # 拼接所有value
+                return " ".join([to_str(v) for v in val.values()])
+            elif isinstance(val, list):
+                return " ".join([to_str(v) for v in val])
+            elif val is None:
+                return ""
+            return str(val)
+
         for item in data:
             if not isinstance(item, dict):
                 continue
-                
-            question = item.get('question', '')
-            answer = item.get('answer', '')
-            context = item.get('context', '')
-            
+            question = to_str(item.get('question', ''))
+            answer = to_str(item.get('answer', ''))
+            context = to_str(item.get('context', ''))
+
             processed_item = {
                 'question': self.clean_text(question),
                 'answer': self.clean_text(answer),
@@ -85,9 +93,7 @@ class Preprocessor:
                 'answer_tokens': self.tokenize(answer),
                 'context_tokens': self.tokenize(context)
             }
-            
             processed.append(processed_item)
-        
         return processed
     
     def create_chunks(self, text: str, chunk_size: int = 512, overlap: int = 50) -> List[str]:

exp_pipeline/pipeline.py

@@ -0,0 +1,50 @@
+"""
+End-to-end pipeline for dataset download, preprocessing, embedding, and indexing.
+"""
+import logging
+from data_processing.data_loader import DataLoader
+from data_processing.preprocessor import Preprocessor
+from retriever.embedder import Embedder
+from retriever.faiss_index import build_faiss_index
+
+logger = logging.getLogger(__name__)
+
+
+def run_pipeline(split: str = "train"):
+    # 1. 下载MS MARCO Passage Ranking数据集
+    data_loader = DataLoader()
+    raw_data = data_loader.get_passage_dataset(split)
+    logger.info(f"Loaded {len(raw_data)} samples from MS MARCO Passage Ranking [{split}]")
+
+    # 2. 预处理数据
+    preprocessor = Preprocessor()
+    # HuggingFace datasets对象转list
+    if hasattr(raw_data, "to_dict"):
+        raw_data = raw_data.to_dict()
+        raw_data = [dict(zip(raw_data.keys(), v)) for v in zip(*raw_data.values())]
+
+    # MS MARCO Passage v2.1: 用passages["passage_text"]字段
+    passages = []
+    for item in raw_data:
+        if "passages" in item and "passage_text" in item["passages"]:
+            passages.extend(item["passages"]["passage_text"])
+    processed = preprocessor.preprocess_passages(passages)
+    texts = [p["text"] for p in processed]
+
+    logger.info(f"Processed {len(texts)} passages")
+
+    # 3. 生产embedding
+    embedder = Embedder(device="cuda")
+    embeddings = embedder.encode(texts)
+    print(f"Embedding shape: {getattr(embeddings, 'shape', None)}")
+    print(f"Texts count: {len(texts)}")
+    if embeddings is None or not hasattr(embeddings, 'shape') or len(embeddings.shape) != 2 or embeddings.shape[0] == 0:
+        raise ValueError("Embeddings is empty or not a 2D array. Check input texts and embedding model.")
+
+    # 4. 建立FAISS索引
+    index = build_faiss_index(embeddings, texts)
+    logger.info("FAISS index built successfully")
+    return index
+
+if __name__ == "__main__":
+    run_pipeline("train")

retriever/faiss_index.py

@@ -1,3 +1,11 @@
+# 工厂函数，供pipeline调用
+def build_faiss_index(embeddings, texts, metadata=None, index_type="IVF"):
+    if embeddings is None or not hasattr(embeddings, 'shape') or len(embeddings.shape) != 2 or embeddings.shape[0] == 0:
+        raise ValueError(f"Embeddings is empty or not a 2D array. Got shape: {getattr(embeddings, 'shape', None)}")
+    dimension = embeddings.shape[1]
+    index = FAISSIndex(dimension, index_type=index_type)
+    index.build_index(embeddings, texts, metadata)
+    return index
 import faiss
 import numpy as np
 import pickle