utils.py

import os
import re
import json
import tiktoken
from pathlib import Path
from datasets import Dataset
from typing import Any, Dict, List
from prompts import INSTRUCTION_PROMPT
from langchain_core.documents import Document
from rapidfireai.automl import RFLangChainRagSpec, RFPromptManager
from langchain_text_splitters import RecursiveCharacterTextSplitter

ANSWER_PATTERN = re.compile(r"<answer>(.*?)</answer>", re.DOTALL)

_ENCODER = tiktoken.get_encoding("gpt2")
def count_tokens(text: str) -> int:
    return len(_ENCODER.encode(text))

CONTEXT_TOKEN_BUDGET = 2000

_USER_PROMPT_TEMPLATE = "Relevant context:\nQuestion:\n"
_DOC_SEPARATOR = "\n\n"
SAFETY_MARGIN = 40 #CHAT_OVERHEAD_TOKENS of 20 + SAFETY_MARGIN of 20

INSTRUCTION_PROMPT_TOKENS = count_tokens(INSTRUCTION_PROMPT)
USER_SCAFFOLD_TOKENS = count_tokens(_USER_PROMPT_TEMPLATE)
DOC_SEPARATOR_TOKENS = count_tokens(_DOC_SEPARATOR)
FIXED_CONTEXT_TOKEN_OVERHEAD = (
    INSTRUCTION_PROMPT_TOKENS
    + USER_SCAFFOLD_TOKENS
    + SAFETY_MARGIN
)


def pack_to_budget(
    docs: List[Document],
    budget: int,
    template,
) -> List[Document]:
    packed: List[Document] = []
    used = 0
    for doc in docs:
        rendered_tokens = count_tokens(template(doc))
        cost = rendered_tokens + (DOC_SEPARATOR_TOKENS if packed else 0)
        if used + cost <= budget:
            packed.append(doc)
            used += cost
    return packed

def LineAwareRecursiveCharacterTextSplitter(chunk_size: int = 512, chunk_overlap: int = 32) -> RecursiveCharacterTextSplitter:
    """
    A text splitter that splits the text into chunks while preserving the line numbers.
    Args:
        chunk_size: The size of the chunks to split the text into. Default is 512.
        chunk_overlap: The overlap between the chunks. Default is 32.

    Returns:
        A RecursiveCharacterTextSplitter object that splits the text into chunks while preserving the line numbers.
    """
    splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        encoding_name="gpt2",
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )
    _orig_split = splitter.split_documents

    def _split_with_lines(documents):
        source_text = {d.metadata["source"]: d.page_content for d in documents}
        chunks = _orig_split(documents)
        cursors: Dict[str, int] = {}
        for c in chunks:
            src = c.metadata["source"]
            text = source_text[src]
            s = text.find(c.page_content, cursors.get(src, 0))
            if s < 0:
                s = cursors.get(src, 0)
            cursors[src] = s
            c.metadata["start_line"] = text.count("\n", 0, s) + 1
            c.metadata["end_line"] = text.count("\n", 0, s + len(c.page_content)) + 1
        return chunks

    splitter.split_documents = _split_with_lines
    return splitter


class MultiResolutionTextSplitter:
    """
    Multi-resolution chunking ensemble: runs N independent
    LineAwareRecursiveCharacterTextSplitter instances at different
    chunk_size / chunk_overlap settings and merges the results into a single
    chunk list, tagging each chunk with its resolution name.

    Letting both small and large chunks coexist in the index lets the
    retriever hit tight ground-truth spans (small chunks) while still
    covering dispersed evidence (large chunks).
    """

    def __init__(
        self,
        resolutions: List[Dict[str, int]] = None,
        deduplicate_at_index: bool = False,
    ):
        if resolutions is None:
            resolutions = [
                {"chunk_size": 128, "chunk_overlap": 32},
                {"chunk_size": 256, "chunk_overlap": 64},
                {"chunk_size": 512, "chunk_overlap": 128},
            ]
        self.resolutions = resolutions
        self.deduplicate_at_index = deduplicate_at_index
        self.splitters: List[tuple[str, RecursiveCharacterTextSplitter]] = []
        for idx, cfg in enumerate(resolutions):
            name = f"res_{idx}_sz{cfg['chunk_size']}"
            splitter = LineAwareRecursiveCharacterTextSplitter(
                chunk_size=cfg["chunk_size"],
                chunk_overlap=cfg["chunk_overlap"],
            )
            self.splitters.append((name, splitter))

    def split_documents(self, documents: List[Document]) -> List[Document]:
        all_chunks: List[Document] = []
        for name, splitter in self.splitters:
            chunks = splitter.split_documents(documents)
            for c in chunks:
                c.metadata["resolution"] = name
            all_chunks.extend(chunks)

        if self.deduplicate_at_index:
            # Same (source, start_line, end_line, content) can arise when
            # a document is shorter than the smallest chunk_size.
            seen: set[tuple] = set()
            deduped: List[Document] = []
            for c in all_chunks:
                key = (
                    c.metadata.get("source"),
                    c.metadata.get("start_line"),
                    c.metadata.get("end_line"),
                    c.page_content,
                )
                if key not in seen:
                    seen.add(key)
                    deduped.append(c)
            all_chunks = deduped

        return all_chunks


def deduplicate_docs_by_span(docs: List[Document]) -> List[Document]:
    """
    Remove exact (source, start_line, end_line) duplicates from a retrieved
    doc list, keeping the first occurrence in rank order. Near-overlapping
    chunks from different resolutions are intentionally preserved — the
    retrieval scorer counts each overlapping span as an independent hit.
    """
    seen: set[tuple] = set()
    out: List[Document] = []
    for d in docs:
        key = (
            d.metadata.get("source"),
            d.metadata.get("start_line"),
            d.metadata.get("end_line"),
        )
        if key not in seen:
            seen.add(key)
            out.append(d)
    return out


def custom_template(doc: Document) -> str:
    return f"{os.path.basename(doc.metadata['source'])}: {doc.page_content}"


def question_preprocess_fn(
    batch: Dict[str, List[Any]], rag: RFLangChainRagSpec, prompt_manager: RFPromptManager
) -> Dict[str, List[Any]]:

    all_context = rag.get_context(batch_queries=batch["question"], serialize=False)
    # Drop exact-span duplicates so identical chunks emitted at different
    # multi-resolution levels don't burn top-k slots.
    all_context = [deduplicate_docs_by_span(docs) for docs in all_context]

    ground_truth_spans = [
        [(item["file"], item["lines"][0], item["lines"][1]) for item in evidence]
        for evidence in batch["source_evidence"]
    ]

    packed_context = [
        pack_to_budget(
            docs,
            budget=CONTEXT_TOKEN_BUDGET - FIXED_CONTEXT_TOKEN_OVERHEAD - count_tokens(question),
            template=custom_template,
        )
        for docs, question in zip(all_context, batch["question"])
    ]

    retrieved_spans = [
        [
            (
                os.path.basename(doc.metadata["source"]),
                doc.metadata["start_line"],
                doc.metadata["end_line"],
            )
            for doc in docs
        ]
        for docs in packed_context
    ]
    
    serialized_context = rag.serialize_documents(packed_context)
    batch["question_id"] = [int(question_id) for question_id in batch["question_id"]]

    return {
        **batch,
        "retrieved_spans": retrieved_spans,
        "ground_truth_spans": ground_truth_spans,
        "serialized_context": serialized_context,
        "prompts": [
            [
                {"role": "system", "content": INSTRUCTION_PROMPT},
                {
                    "role": "user",
                    "content": f"Relevant context:\n{context}\nQuestion:\n{question}",
                },
            ]
            for question, context in zip(batch["question"], serialized_context)
        ],
    }


def response_postprocess_fn(batch: Dict[str, List[Any]]) -> Dict[str, List[Any]]:
    extracted = []
    for text in batch["generated_text"]:
        matches = ANSWER_PATTERN.findall(text)
        extracted.append(matches[-1].strip() if matches else "")
    batch["generated_text"] = extracted
    return batch

def load_eval_dataset(eval_dataset: Path) -> Dataset:
    with open(eval_dataset, "r") as f:
        data = json.load(f)
    return Dataset.from_list(data)