{"slug": "building-a-rag-pipeline-with-llm-embeddings-and-scikit-learn", "title": "Building a RAG Pipeline with LLM Embeddings and Scikit-learn", "summary": "A tutorial published on Machine Learning Mastery details how to build a lightweight, in-memory retrieval augmented generation (RAG) pipeline using scikit-learn's NearestNeighbors with cosine distance for vector retrieval, sentence-transformers for local embeddings, and a lightweight open-source Hugging Face LLM for generation, avoiding paid APIs and managed vector databases. The guide uses a nine-document dataset of Asian country descriptions hosted in a public GitHub repository and recommends running the pipeline on Google Colab's free T4 GPU.", "body_md": "In this article, you will learn how to build a lightweight, end-to-end RAG pipeline using scikit-learn for vector retrieval and open-source models for embeddings and text generation, without the use of paid APIs or heavyweight frameworks.\n\nTopics we will cover include:\n\n- How to download a document dataset, chunk it into semantic pieces, and generate embeddings using a local sentence-transformer model.\n- How to build a custom retriever class using scikit-learn’s NearestNeighbors with cosine distance as a stand-in for a vector database.\n- How to wire the retriever together with a lightweight open-source LLM to produce a fully functional, in-memory RAG pipeline.\n\n## Introduction\n\nModern retrieval augmented generation (**RAG systems**) are often thought of as inseparable from an expansive ecosystem of frameworks, including typically expensive, managed vector databases to store document collections for retrieval.\n\nAt the core, however, RAG systems rely on a fairly simple notion: calculating distances between semantic vectors representing texts. It’s not strictly necessary to navigate the intricacies of a dedicated AI infrastructure when working with a small-to-medium-sized corpus, prototyping a new proof-of-concept, or deploying to environments under cost constraints. Instead, we can use **scikit-learn** procedures together with **open-source LLMs** to build transparent, fast RAG pipelines that remain fully functional. This article shows how.\n\n## Step-by-Step Guide\n\nThe solution provided below is completely free to use, relying on `sentence-transformers` to build local embeddings from a text dataset and a lightweight pre-trained LLM from Hugging Face for the generation stage. If you need to brush up on basic RAG concepts, check out [this article](https://machinelearningmastery.com/understanding-rag-part-ii-how-classic-rag-works/) first.\n\nWhile the pipeline can be implemented and run in a local IDE, for demonstration purposes I recommend running the entire pipeline on Google Colab, using its free T4 GPU for more efficient execution. If you follow this approach, start by installing and importing the required libraries and dependencies:\n\n``` python\n!pip install -q sentence-transformers transformers\n\nimport os\nimport glob\nimport urllib.request\nimport zipfile\nimport numpy as np\nimport pandas as pd\nfrom sklearn.neighbors import NearestNeighbors\nfrom sentence_transformers import SentenceTransformer\nfrom transformers import pipeline\n\nprint(\"Libraries loaded successfully!\")\n\n12345678910111213\n\n!pip install -q sentence-transformers transformers import osimport globimport urllib.requestimport zipfileimport numpy as npimport pandas as pdfrom sklearn.neighbors import NearestNeighborsfrom sentence_transformers import SentenceTransformerfrom transformers import pipeline print(\"Libraries loaded successfully!\")\n```\n\nI made available a small **dataset** containing nine text files — brief descriptions of Asian countries — in a public GitHub repository. You can download it, extract the files, and save the texts in a simple list as follows:\n\n```\n# Downloading and extracting the toy dataset, containing 9 documents\nurl = \"https://github.com/gakudo-ai/open-datasets/blob/main/asia_documents.zip?raw=true\"\nzip_path = \"asia_documents.zip\"\nextract_dir = \"asia_documents\"\n\nurllib.request.urlretrieve(url, zip_path)\nwith zipfile.ZipFile(zip_path, 'r') as zip_ref:\n    zip_ref.extractall(extract_dir)\n\n# Reading all extracted text files\ndocuments = []\nfor filepath in glob.glob(f\"{extract_dir}/**/*.txt\", recursive=True):\n    with open(filepath, 'r', encoding='utf-8') as file:\n        documents.append(file.read())\n\nprint(f\"Loaded {len(documents)} documents.\")\n# If the zip contained a CSV instead, you'd use pd.read_csv() here.\n\n1234567891011121314151617\n\n# Downloading and extracting the toy dataset, containing 9 documentsurl = \"https://github.com/gakudo-ai/open-datasets/blob/main/asia_documents.zip?raw=true\"zip_path = \"asia_documents.zip\"extract_dir = \"asia_documents\" urllib.request.urlretrieve(url, zip_path)with zipfile.ZipFile(zip_path, 'r') as zip_ref:    zip_ref.extractall(extract_dir) # Reading all extracted text filesdocuments = []for filepath in glob.glob(f\"{extract_dir}/**/*.txt\", recursive=True):    with open(filepath, 'r', encoding='utf-8') as file:        documents.append(file.read()) print(f\"Loaded {len(documents)} documents.\")# If the zip contained a CSV instead, you'd use pd.read_csv() here.\n```\n\nThe dataset we just loaded will constitute our miniature document base for RAG. We need to put it in the right format — that is, turn the texts into embedding vectors. Before doing that, it is common practice in RAG to **chunk documents into smaller, semantic pieces** that the embedding model — responsible for mapping pieces of text into numerical embeddings — can process more effectively. The texts in our dataset don’t really require chunking because they are quite short, but we are simulating this step to illustrate best practices in real-world RAG pipelines.\n\nWe define a native chunking function, using smaller-than-usual values for chunk size and overlap (both measured in words) so that we can observe some actual effect on our small documents:\n\n``` python\ndef chunk_text(text, chunk_size=25, overlap=5):\n    \"\"\"Splits text into chunks natively, without external frameworks.\"\"\"\n    words = text.split()\n    chunks = []\n    for i in range(0, len(words), chunk_size - overlap):\n        chunk = \" \".join(words[i:i + chunk_size])\n        chunks.append(chunk)\n    return chunks\n\n# Processing all documents into a flat list of chunks\nall_chunks = []\nfor doc in documents:\n    all_chunks.extend(chunk_text(doc))\n\nprint(f\"Created {len(all_chunks)} chunks from the original documents.\")\n\n123456789101112131415\n\ndef chunk_text(text, chunk_size=25, overlap=5):    \"\"\"Splits text into chunks natively, without external frameworks.\"\"\"    words = text.split()    chunks = []    for i in range(0, len(words), chunk_size - overlap):        chunk = \" \".join(words[i:i + chunk_size])        chunks.append(chunk)    return chunks # Processing all documents into a flat list of chunksall_chunks = []for doc in documents:    all_chunks.extend(chunk_text(doc)) print(f\"Created {len(all_chunks)} chunks from the original documents.\")\n```\n\nOutput:\n\n```\nCreated 31 chunks from the original documents.\n\n1\n\nCreated 31 chunks from the original documents.\n```\n\nIn a real-world application, you will likely use significantly larger values for both arguments in the chunking function.\n\nNext, we **generate the embeddings** using a popular lightweight model for this task: the `'all-MiniLM-L6-v2'` **sentence-transformer**:\n\n```\n# Initializing the embedding model (runs locally on Colab's free CPU/GPU)\nembedding_model = SentenceTransformer('all-MiniLM-L6-v2')\n\n# Generating embeddings for all text chunks\nprint(\"Generating embeddings... (this is fast on Colab!)\")\nchunk_embeddings = embedding_model.encode(all_chunks, show_progress_bar=True)\n\n123456\n\n# Initializing the embedding model (runs locally on Colab's free CPU/GPU)embedding_model = SentenceTransformer('all-MiniLM-L6-v2') # Generating embeddings for all text chunksprint(\"Generating embeddings... (this is fast on Colab!)\")chunk_embeddings = embedding_model.encode(all_chunks, show_progress_bar=True)\n```\n\nOnce we have a collection of embeddings, it’s time to define one of the key components of any RAG pipeline: the **retriever**. Since we are doing this the scikit-learn way, we will wrap `sklearn.neighbors.NearestNeighbors` — a simple yet effective distance calculation algorithm — into a clean, reusable class that mimics the functionality of a vector database, all based on cosine distance, which is ideal for embedding vector representations.\n\n``` python\nclass SklearnRetriever:\n    def __init__(self, n_neighbors=3):\n        # Using 'cosine' metric because LLM embeddings measure semantic similarity based on angles between vectors\n        self.retriever = NearestNeighbors(n_neighbors=n_neighbors, metric='cosine', algorithm='brute')\n        self.chunks = []\n\n    def fit(self, chunks, embeddings):\n        \"\"\"Indexes the documents.\"\"\"\n        self.chunks = chunks\n        self.retriever.fit(embeddings)\n        print(\"Vector index built successfully!\")\n\n    def retrieve(self, query_embedding):\n        \"\"\"Returns the most relevant chunks for a given query.\"\"\"\n        distances, indices = self.retriever.kneighbors([query_embedding])\n        \n        # Fetching the actual text strings using the retrieved indices\n        results = [self.chunks[idx] for idx in indices[0]]\n        return results\n\n# Instantiating and fitting our custom retriever\nvector_db = SklearnRetriever(n_neighbors=3)\nvector_db.fit(all_chunks, chunk_embeddings)\n\n1234567891011121314151617181920212223\n\nclass SklearnRetriever:    def __init__(self, n_neighbors=3):        # Using 'cosine' metric because LLM embeddings measure semantic similarity based on angles between vectors        self.retriever = NearestNeighbors(n_neighbors=n_neighbors, metric='cosine', algorithm='brute')        self.chunks = []     def fit(self, chunks, embeddings):        \"\"\"Indexes the documents.\"\"\"        self.chunks = chunks        self.retriever.fit(embeddings)        print(\"Vector index built successfully!\")     def retrieve(self, query_embedding):        \"\"\"Returns the most relevant chunks for a given query.\"\"\"        distances, indices = self.retriever.kneighbors([query_embedding])                # Fetching the actual text strings using the retrieved indices        results = [self.chunks[idx] for idx in indices[0]]        return results # Instantiating and fitting our custom retrievervector_db = SklearnRetriever(n_neighbors=3)vector_db.fit(all_chunks, chunk_embeddings)\n```\n\nThe other major component of the pipeline is the **generator**. Here too we keep things lightweight, choosing a small, open-source LLM trained for natural language generation: `Qwen/Qwen2.5-0.5B-Instruct`. This model is small enough to run smoothly on a Colab notebook without risking a RAM crash.\n\n```\n# Loading a lightweight causal LLM built for the \"text-generation\" pipeline\nllm = pipeline(\"text-generation\", model=\"Qwen/Qwen2.5-0.5B-Instruct\", device_map=\"auto\")\n\ndef generate_rag_answer(query, retrieved_context):\n    \"\"\"Combines the prompt, context, and query for the LLM using chat formatting.\"\"\"\n    context_str = \" --- \".join(retrieved_context)\n    \n    messages = [\n        {\"role\": \"system\", \"content\": \"You are a helpful assistant. Answer the user's question using ONLY the provided context. If the answer is not in the context, say 'I cannot find the answer in the provided documents.'\"},\n        {\"role\": \"user\", \"content\": f\"Context:\\n{context_str}\\n\\nQuestion: {query}\"}\n    ]\n    \n    # Apply the model's chat template\n    prompt = llm.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n    \n    # Generating the response\n    response = llm(prompt, max_new_tokens=150, do_sample=False, temperature=0.1)\n    \n    # Extracting only the newly generated text past the prompt\n    generated_text = response[0]['generated_text']\n    answer = generated_text[len(prompt):].strip()\n    return answer\n\n12345678910111213141516171819202122\n\n# Loading a lightweight causal LLM built for the \"text-generation\" pipelinellm = pipeline(\"text-generation\", model=\"Qwen/Qwen2.5-0.5B-Instruct\", device_map=\"auto\") def generate_rag_answer(query, retrieved_context):    \"\"\"Combines the prompt, context, and query for the LLM using chat formatting.\"\"\"    context_str = \" --- \".join(retrieved_context)        messages = [        {\"role\": \"system\", \"content\": \"You are a helpful assistant. Answer the user's question using ONLY the provided context. If the answer is not in the context, say 'I cannot find the answer in the provided documents.'\"},        {\"role\": \"user\", \"content\": f\"Context:\\n{context_str}\\n\\nQuestion: {query}\"}    ]        # Apply the model's chat template    prompt = llm.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)        # Generating the response    response = llm(prompt, max_new_tokens=150, do_sample=False, temperature=0.1)        # Extracting only the newly generated text past the prompt    generated_text = response[0]['generated_text']    answer = generated_text[len(prompt):].strip()    return answer\n```\n\nAll the pieces of the pipeline are now in place. The last step is to wire them together and **try it out** by asking our system a question about its knowledge base (Asian countries):\n\n``` python\ndef ask_rag_pipeline(question):\n    print(f\"Question: {question}\")\n    \n    # 1. Embedding the query\n    query_vector = embedding_model.encode(question)\n    \n    # 2. Retrieving context via Scikit-Learn\n    relevant_chunks = vector_db.retrieve(query_vector)\n    \n    # 3. Generating Answer\n    answer = generate_rag_answer(question, relevant_chunks)\n    \n    print(f\"\\nRAG Answer: {answer}\")\n    print(\"\\n--- Sources used ---\")\n    for i, chunk in enumerate(relevant_chunks):\n        print(f\"{i+1}. {chunk[:100]}...\") # Truncated for readability\n\n# Testing it out\ntest_question = \"Which Asian countries are best known for islands and beaches?\"\nask_rag_pipeline(test_question)\n\n1234567891011121314151617181920\n\ndef ask_rag_pipeline(question):    print(f\"Question: {question}\")        # 1. Embedding the query    query_vector = embedding_model.encode(question)        # 2. Retrieving context via Scikit-Learn    relevant_chunks = vector_db.retrieve(query_vector)        # 3. Generating Answer    answer = generate_rag_answer(question, relevant_chunks)        print(f\"\\nRAG Answer: {answer}\")    print(\"\\n--- Sources used ---\")    for i, chunk in enumerate(relevant_chunks):        print(f\"{i+1}. {chunk[:100]}...\") # Truncated for readability # Testing it outtest_question = \"Which Asian countries are best known for islands and beaches?\"ask_rag_pipeline(test_question)\n```\n\nOutput:\n\n```\nQuestion: Which Asian countries are best known for islands and beaches?\n\nRAG Answer: The Asian countries that are best known for islands and beaches include:\n\n1. Philippines - Known for its over 7,000 islands and vibrant festivals.\n2. Thailand - Famous for its tropical beaches and bustling street food culture.\n\nThese two countries have a significant number of islands and offer both beautiful natural scenery and diverse beach experiences.\n\n--- Sources used ---\n1. mountainous landscapes and cultural cities like Chiang Mai, while the south offers world-renowned is...\n2. The Philippines is an archipelago of over 7,000 islands in Southeast Asia, known for its stunning be...\n3. Thailand is a Southeast Asian country famous for its tropical beaches, ornate temples, and bustling ...\n\n12345678910111213\n\nQuestion: Which Asian countries are best known for islands and beaches? RAG Answer: The Asian countries that are best known for islands and beaches include: 1. Philippines - Known for its over 7,000 islands and vibrant festivals.2. Thailand - Famous for its tropical beaches and bustling street food culture. These two countries have a significant number of islands and offer both beautiful natural scenery and diverse beach experiences. --- Sources used ---1. mountainous landscapes and cultural cities like Chiang Mai, while the south offers world-renowned is...2. The Philippines is an archipelago of over 7,000 islands in Southeast Asia, known for its stunning be...3. Thailand is a Southeast Asian country famous for its tropical beaches, ornate temples, and bustling ...\n```\n\nTry asking different questions by replacing the text in `test_question`. The behavior and outputs are reasonable given the small-scale setup we just built — though don’t expect commercial LLM-level results: we are combining a small model with a tiny knowledge base.\n\n## Final Remarks\n\nIn this article, we built a lightweight, end-to-end RAG pipeline that avoids the overhead of heavyweight frameworks or paid APIs. It relies solely on scikit-learn vector retrieval and open-source models for both embeddings and generation, producing a transparent, in-memory system. This approach makes a solid starting point for prototyping, experimenting with different retrieval strategies, and validating concepts before scaling up.", "url": "https://wpnews.pro/news/building-a-rag-pipeline-with-llm-embeddings-and-scikit-learn", "canonical_source": "https://machinelearningmastery.com/building-a-rag-pipeline-with-llm-embeddings-and-scikit-learn/", "published_at": "2026-10-13 11:00:00+00:00", "updated_at": "2026-10-15 05:05:39.732165+00:00", "lang": "en", "topics": ["artificial-intelligence", "large-language-models", "generative-ai", "ai-tools", "developer-tools"], "entities": ["scikit-learn", "sentence-transformers", "Hugging Face", "Google Colab", "Machine Learning Mastery", "NearestNeighbors", "GitHub"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/building-a-rag-pipeline-with-llm-embeddings-and-scikit-learn", "markdown": "https://wpnews.pro/news/building-a-rag-pipeline-with-llm-embeddings-and-scikit-learn.md", "text": "https://wpnews.pro/news/building-a-rag-pipeline-with-llm-embeddings-and-scikit-learn.txt", "jsonld": "https://wpnews.pro/news/building-a-rag-pipeline-with-llm-embeddings-and-scikit-learn.jsonld"}}