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()