first commit

config.yaml

@@ -0,0 +1,37 @@
+# Pipeline Configuration for ArXiv RAG
+project_name: "arxiv_cyber_advisor"
+
+# Stage 1: Data Acquisition
+data_ingestion:
+  category: "cs.AI"
+  limit: 20
+  save_local: true
+  raw_data_path: "data/raw_arxiv.csv"
+
+# Stage 2: Processing & Embedding
+embedding:
+  model_name: "all-MiniLM-L6-v2"
+  device: "cpu"  # Change to "cuda" if testing on a GPU machine
+
+chunking:
+  technique: "recursive"
+  chunk_size: 500
+  chunk_overlap: 50
+
+# Stage 3: Vector Database (Pinecone)
+vector_db:
+  index_name: "arxiv-index"
+  dimension: 384
+  metric: "cosine"
+
+# Stage 4: Retrieval & Re-ranking
+retrieval:
+  top_k_hybrid: 10
+  rerank_model: "cross-encoder/ms-marco-MiniLM-L-6-v2"
+  top_k_final: 3
+
+# Stage 5: Generation (LLM)
+llm:
+  model_id: "meta-llama/Meta-Llama-3-8B-Instruct"
+  max_new_tokens: 500
+  temperature: 0.1

data_processor.py

@@ -0,0 +1,69 @@
+import arxiv
+import pandas as pd
+from langchain_text_splitters import RecursiveCharacterTextSplitter, CharacterTextSplitter
+from sentence_transformers import SentenceTransformer
+
+def fetch_arxiv_data(category="cs.AI", limit=10):
+    """Fetches paper metadata and abstracts from arXiv."""
+    client = arxiv.Client()
+    search = arxiv.Search(
+        query=f"cat:{category}",
+        max_results=limit,
+        sort_by=arxiv.SortCriterion.SubmittedDate
+    )
+    
+    results = []
+    for r in client.results(search):
+        results.append({
+            "id": r.entry_id.split('/')[-1],
+            "title": r.title,
+            "abstract": r.summary.replace('\n', ' '),
+            "url": r.pdf_url
+        })
+    
+    return pd.DataFrame(results)
+
+def get_text_splitter(technique="recursive", chunk_size=500, chunk_overlap=50):
+    """Returns a splitter based on the chosen technique and parameters."""
+    if technique == "recursive":
+        return RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size, 
+            chunk_overlap=chunk_overlap
+        )
+    elif technique == "character":
+        return CharacterTextSplitter(
+            separator="\n",
+            chunk_size=chunk_size, 
+            chunk_overlap=chunk_overlap
+        )
+    else:
+        raise ValueError(f"Technique '{technique}' not supported.")
+
+def process_to_chunks(df, model_name='all-MiniLM-L6-v2', technique="recursive", chunk_size=500, chunk_overlap=50):
+    """Splits abstracts into chunks and generates embeddings with custom parameters."""
+    
+    # Initialize the specific model requested
+    print(f"🔧 Initializing Model: {model_name}...")
+    model = SentenceTransformer(model_name)
+    
+    # Initialize the specific splitter requested
+    splitter = get_text_splitter(technique, chunk_size, chunk_overlap)
+    
+    processed_chunks = []
+    
+    for _, row in df.iterrows():
+        chunks = splitter.split_text(row['abstract'])
+        
+        for i, text in enumerate(chunks):
+            embedding = model.encode(text).tolist()
+            
+            processed_chunks.append({
+                "id": f"{row['id']}-chunk-{i}",
+                "values": embedding,
+                "metadata": {
+                    "title": row['title'],
+                    "text": text,
+                    "url": row['url']
+                }
+            })
+    return processed_chunks

main.py

@@ -0,0 +1,28 @@
+import yaml
+from data_processor import fetch_arxiv_data, process_to_chunks
+
+def load_config():
+    with open("config.yaml", "r") as f:
+        return yaml.safe_load(f)
+
+def main():
+    config = load_config()
+    
+    # Run Stage 1
+    raw_data = fetch_arxiv_data(
+        category=config['data_ingestion']['category'], 
+        limit=config['data_ingestion']['limit']
+    )
+
+    # Run Stage 2 using YAML defaults
+    final_chunks = process_to_chunks(
+        raw_data,
+        model_name=config['embedding']['model_name'],
+        chunk_size=config['chunking']['chunk_size'],
+        chunk_overlap=config['chunking']['chunk_overlap']
+    )
+
+    print(f"✅ Pipeline finished with {len(final_chunks)} chunks.")
+
+if __name__ == "__main__":
+    main()

retriever.py

@@ -0,0 +1,78 @@
+import numpy as np
+from rank_bm25 import BM25Okapi
+from sentence_transformers import CrossEncoder
+
+class HybridRetriever:
+    def __init__(self, final_chunks, embed_model, rerank_model_name='cross-encoder/ms-marco-MiniLM-L-6-v2'):
+        """
+        Initializes the search engines.
+        :param final_chunks: The list of chunk dictionaries with metadata.
+        :param embed_model: The SentenceTransformer model object used for query embedding.
+        """
+        self.final_chunks = final_chunks
+        self.embed_model = embed_model
+        self.rerank_model = CrossEncoder(rerank_model_name)
+        
+        # Initialize BM25 locally
+        self.tokenized_corpus = [chunk['metadata']['text'].lower().split() for chunk in final_chunks]
+        self.bm25 = BM25Okapi(self.tokenized_corpus)
+
+    def _rrf_score(self, semantic_results, bm25_results, k=60):
+        """
+        Reciprocal Rank Fusion (RRF) Implementation.
+        Score = 1 / (k + rank)
+        """
+        scores = {}
+        
+        # Rank is index + 1
+        for rank, chunk in enumerate(semantic_results):
+            scores[chunk] = scores.get(chunk, 0) + 1 / (k + rank + 1)
+            
+        for rank, chunk in enumerate(bm25_results):
+            scores[chunk] = scores.get(chunk, 0) + 1 / (k + rank + 1)
+            
+        # Sort by score descending
+        sorted_chunks = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+        return [item[0] for item in sorted_chunks]
+
+    def search(self, query, index, top_k=10, mode="all", rerank_type="cross-encoder"):
+        """
+        :param mode: "semantic", "bm25", or "all"
+        :param rerank_type: "cross-encoder", "rrf", or "none"
+        """
+        semantic_chunks = []
+        bm25_chunks = []
+
+        # A. Semantic Search
+        if mode in ["semantic", "all"]:
+            query_vector = self.embed_model.encode(query).tolist()
+            res = index.query(vector=query_vector, top_k=top_k, include_metadata=True)
+            semantic_chunks = [match['metadata']['text'] for match in res['matches']]
+
+        # B. Keyword Search
+        if mode in ["bm25", "all"]:
+            tokenized_query = query.lower().split()
+            bm25_scores = self.bm25.get_scores(tokenized_query)
+            top_indices = np.argsort(bm25_scores)[::-1][:top_k]
+            bm25_chunks = [self.final_chunks[i]['metadata']['text'] for i in top_indices]
+
+        # C. Combination and Re-Ranking
+        if mode == "semantic":
+            combined = semantic_chunks
+        elif mode == "bm25":
+            combined = bm25_chunks
+        else:
+            # Mode is "all"
+            if rerank_type == "rrf":
+                return self._rrf_score(semantic_chunks, bm25_chunks)[:3]
+            else:
+                combined = list(set(semantic_chunks + bm25_chunks))
+
+        # D. Cross-Encoder Re-Ranking
+        if rerank_type == "cross-encoder" and len(combined) > 0:
+            pairs = [[query, chunk] for chunk in combined]
+            scores = self.rerank_model.predict(pairs)
+            results = sorted(zip(combined, scores), key=lambda x: x[1], reverse=True)
+            return [res[0] for res in results[:3]]
+
+        return combined[:3]

vector_db.py

@@ -0,0 +1,33 @@
+import time
+from pinecone import Pinecone, ServerlessSpec
+
+def get_pinecone_index(api_key, index_name, dimension=384, metric="cosine"):
+    """Initializes Pinecone and returns the index object, creating it if necessary."""
+    pc = Pinecone(api_key=api_key)
+    
+    # Check if index exists
+    existing_indexes = [idx.name for idx in pc.list_indexes()]
+    
+    if index_name not in existing_indexes:
+        print(f"Creating new Pinecone index: {index_name}...")
+        pc.create_index(
+            name=index_name,
+            dimension=dimension,
+            metric=metric,
+            spec=ServerlessSpec(cloud="aws", region="us-east-1")
+        )
+        # Wait for index to be ready
+        while not pc.describe_index(index_name).status['ready']:
+            time.sleep(1)
+    
+    return pc.Index(index_name)
+
+def upsert_to_pinecone(index, chunks, batch_size=100):
+    """Upserts chunks to Pinecone in manageable batches."""
+    print(f"Uploading {len(chunks)} chunks to Pinecone...")
+    
+    for i in range(0, len(chunks), batch_size):
+        batch = chunks[i : i + batch_size]
+        index.upsert(vectors=batch)
+        
+    print("✅ Upsert complete.")