boardgame_rag.py

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
from pathlib import Path
from collections import Counter
import ast
import json
import shutil
import time
from textwrap import dedent

import numpy as np
import pandas as pd

from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter, MarkdownHeaderTextSplitter
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings

class BoardGameRag():

    def __init__(self, model_name:str, tokenizer_name: str | None, children:list[Document] | None=None, parents:list[Document] | None=None, max_new_tokens = 500, do_sample = False, persist_directory: str|None=None):
        if tokenizer_name is None:
            tokenizer_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map = "auto", dtype = torch.bfloat16)
        self.max_new_tokens = max_new_tokens
        self.do_sample = do_sample
        self.pipe = pipeline(
            "text-generation",
            model = self.model,
            tokenizer = self.tokenizer,
            torch_dtype = torch.bfloat16,
            device_map = "auto",
            max_new_tokens = self.max_new_tokens,
            do_sample = self.do_sample
        )
        self.children = children
        self.parents = parents
        self.embedding_model = "Qwen/Qwen3-Embedding-4B"
        self.retrieval_enabled = children is not None and parents is not None
        self.persist_directory = persist_directory

        if self.retrieval_enabled:
            self.parent_lookup = self._build_parent_lookup()

            self.embeddings = HuggingFaceEmbeddings(
                model_name = self.embedding_model,
                model_kwargs = {"device" : "cuda" if torch.cuda.is_available() else "cpu" },
                encode_kwargs={"normalize_embeddings": True}
            )

            if self.persist_directory is not None and Path(self.persist_directory).exists():
                self.vector_store = Chroma(
                    embedding_function = self.embeddings,
                    persist_directory = self.persist_directory
                )
            else:    
                self.vector_store = Chroma.from_documents(
                    documents = self.children, 
                    embedding = self.embeddings, 
                    ids = [child.metadata["chunk_id"] for child in self.children],
                    persist_directory = self.persist_directory
                    )
        else:
            self.parent_lookup = []
            self.embeddings = None
            self.vector_store = None

    def _build_parent_lookup(self):
        parent_lookup = {}
        for parent in self.parents:
            parent_id = (
                parent.metadata["game"],
                parent.metadata["section_id"],
                parent.metadata["subsection_id"]
            )
            parent_lookup[parent_id] = parent
        return parent_lookup

    def _retrieve_child_chunks(self, question):

        search_kwargs={"k": 5}
        
        retriever = self.vector_store.as_retriever(search_type = "similarity",  search_kwargs = search_kwargs)

        retrieved_chunks = retriever.invoke(question)

        return retrieved_chunks

    def _expand_children_to_parents(self, child_chunks):

        parent_documents = []
        seen_parent_ids = set()

        for child in child_chunks:
            parent_id = (
                child.metadata["game"],
                child.metadata["section_id"],
                child.metadata["subsection_id"]
            )

            if parent_id not in seen_parent_ids:
                parent = self.parent_lookup.get(parent_id)

                if parent is None:
                    raise KeyError(
                        f"No parent subsection found for {parent_id}"
                    )
                parent_documents.append(parent)
                seen_parent_ids.add(parent_id)

        return parent_documents

    def retrieve(self, question):
        
        if not self.retrieval_enabled:
            raise RuntimeError(
            "This BoardGameRag instance was created without `documents`/`subsections`, "
            "so it has no vector store to retrieve from. Either construct it with both, "
            "or use generate_from_context() with pre-supplied contexts instead."
        )
        
        child_chunks = self._retrieve_child_chunks(question = question)

        parent_documents = self._expand_children_to_parents(child_chunks)

        return {
            "child_chunks" : child_chunks,
            "parent_documents" : parent_documents
        }
    
    def _format_parent_context(self, parent_documents: list[Document]):
        formatted_parents = []

        for parent in parent_documents:
            metadata = parent.metadata

            citation_label = f"{metadata['section_id']}_{metadata['subsection_id']}"
            formatted_parents.append(
                f"{citation_label}:\n"
                f"{parent.page_content.strip()}"
                )
        return "\n\n".join(formatted_parents)
    
    def _build_rag_prompt(self, question, context, prompt: str | None=None):

        if prompt is None:
            prompt = dedent(
            f"""
            You are a helpful assistant answering questions about how to play board games. 
            Answer the user's question thoroughly using only the provided context. 
            If the answer is not directly supported by the context you must say that you cannot answer the question with the information provided. 
            If the answer differs across game modes, variants, or optional rules described in the context, summarize the differences and specify which version each applies to. 
            Do not provide gameplay tips or strategies unless they are explicitly mentioned in the context.
            """).strip()

        return dedent(f"""
        {prompt}

        CONTEXT
        {context}

        QUESTION
        {question}

        ANSWER:
        """).strip()
    
    def generate_answer(self, question, prompt):
        retrieval = self.retrieve(question)
        context = self._format_parent_context(retrieval["parent_documents"])
        prompt = self._build_rag_prompt(question, context, prompt)

        messages = [
            {
                "role": "user",
                "content": prompt,
            }
        ]

        prompt = self.tokenizer.apply_chat_template(
            messages,
            tokenize = False,
            add_generation_prompt = True
        )

        output = self.pipe(prompt, return_full_text = False)

        answer = output[0]['generated_text'].strip()

        return {
        "answer": answer,
        "child_documents": retrieval["child_chunks"],
        "parent_documents": retrieval["parent_documents"],
        "context": context,
        "prompt": prompt
    }

    def generate_from_context(self, question: str, contexts: list[str], prompt: str | None=None):
        context = "\n\n".join(contexts)

        formatted_prompt = self._build_rag_prompt(
            question=question,
            context=context,
            prompt=prompt,
        )

        messages = [
            {
                "role" : "user",
                "content" : formatted_prompt
            }
        ]

        prompt = self.tokenizer.apply_chat_template(
            messages,
            tokenize = False,
            add_generation_prompt = True
        )
        output = self.pipe(prompt, return_full_text = False)

        generated_text = output[0]["generated_text"]

        return generated_text.strip()