ragbits-evaluate 0.0.30.dev29302392__py3-none-any.whl

ragbits/evaluate/dataset_generator/pipeline.py

@@ -0,0 +1,141 @@
+import sys
+from typing import Any
+
+from datasets import Dataset
+from distilabel.pipeline import Pipeline
+from distilabel.steps.base import Step
+from omegaconf import DictConfig, OmegaConf
+from pydantic import BaseModel
+
+from ragbits.core.utils.config_handling import import_by_path
+
+module = sys.modules[__name__]
+
+
+class LLMConfigForTask(BaseModel):
+    """
+    Configuration for the LLM (Language Model) associated with a specific task.
+
+    Attributes:
+        provider_type (str): The type of LLM provider.
+        kwargs (dict): Additional parameters or settings for the LLM provider.
+    """
+
+    provider_type: str
+    kwargs: dict
+
+
+class TaskConfig(BaseModel):
+    """
+    Configuration for an individual task in the dataset generation pipeline.
+
+    Attributes:
+        type: str: type of the task
+        llm (LLMConfigForTask): The configuration for the LLM used in this task.
+        kwargs (dicts): Optional additional parameters or settings for the task.
+        filters (list[str] | None): Optional filters to apply during the task. Defaults to None.
+    """
+
+    type: str
+    llm: LLMConfigForTask
+    kwargs: dict | None = None
+    filters: list[str] | None = None
+
+
+class DatasetGenerationPipelineConfig(BaseModel):
+    """
+    Configuration for the entire dataset generation pipeline.
+
+    Attributes:
+        name (str): The name of the dataset generation pipeline.
+        input_name (str): The name of the input resource or dataset.
+        tasks (list[TaskConfig]): A list of task configurations included in the pipeline.
+    """
+
+    name: str
+    input_name: str
+    tasks: list[TaskConfig]
+
+    @classmethod
+    def from_dict_config(cls, dict_config: DictConfig) -> "DatasetGenerationPipelineConfig":
+        """
+        Creates an instance of `DatasetGenerationPipelineConfig` from a dictionary-based configuration.
+
+        Args:
+            dict_config (DictConfig): A configuration object containing pipeline details.
+
+        Returns:
+            DatasetGenerationPipelineConfig: An instance populated with data from the given configuration.
+
+        """
+        name = dict_config.name
+        input_name = dict_config.input_name
+        tasks = [
+            TaskConfig(
+                type=task_config.type,
+                llm=LLMConfigForTask(
+                    provider_type=task_config.llm.provider_type,
+                    kwargs=OmegaConf.to_container(task_config.llm.kwargs),  # type: ignore
+                ),
+                kwargs=OmegaConf.to_container(task_config.kwargs),  # type: ignore
+                filters=getattr(task_config, "filters", None),
+            )
+            for task_config in dict_config.tasks
+        ]
+        return cls(name=name, input_name=input_name, tasks=tasks)
+
+
+class DatasetGenerationPipeline:
+    """A pipeline for dataset generation"""
+
+    def __init__(self, config: DatasetGenerationPipelineConfig):
+        self.config = config
+        self._instantiate_pipeline()
+
+    @classmethod
+    def from_dict_config(cls, dict_config: DictConfig) -> "DatasetGenerationPipeline":
+        """
+        Instantiates the pipeline from dict config validated through pydantic base model
+        Returns:
+            DatasetGenerationPipeline
+        """
+        config = DatasetGenerationPipelineConfig.from_dict_config(dict_config=dict_config)
+        return cls(config=config)
+
+    def __call__(self, corpus: list[str]) -> Dataset:
+        """
+        Generates a dataset from a corpus or list of topics
+        Args:
+            corpus: a corpus of information or list of topics
+        Returns:
+            dataset instance
+        """
+        dataset = Dataset.from_dict({self.config.input_name: corpus})
+        distiset = self.pipeline.run(use_cache=False, dataset=dataset)
+        result = distiset["default"]["train"]
+        result = result.remove_columns(["distilabel_metadata", "model_name"])
+        return result
+
+    def _parse_pipeline_steps(self) -> list[Step]:
+        tasks = []
+        for task_config in self.config.tasks:
+            llm_config = task_config.llm
+            llm = import_by_path(llm_config.provider_type, module)(**llm_config.kwargs)
+            task_kwargs: dict[Any, Any] = {"llm": llm}
+            task_kwargs.update(task_config.kwargs or {})  # type: ignore
+            task = import_by_path(task_config.type, module)(**task_kwargs)
+            tasks.append(task)
+            filter_types = getattr(task_config, "filters", None) or []
+            for filter_type in filter_types:
+                filter = import_by_path(filter_type, module)(tasks[-1])
+                tasks.append(filter)
+        return tasks
+
+    def _instantiate_pipeline(self) -> None:
+        with Pipeline(self.config.name) as self.pipeline:
+            tasks = self._parse_pipeline_steps()
+            prev_task = None
+            for task in tasks:
+                if prev_task:
+                    prev_task >> task
+                prev_task = task

ragbits/evaluate/dataset_generator/prompts/corpus_generation.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel
+
+from ragbits.core.prompt import Prompt
+
+
+class BasicCorpusGenerationPromptInput(BaseModel):
+    """A definition of input for corpus generation task"""
+
+    query: str
+
+
+class BasicCorpusGenerationPrompt(Prompt[BasicCorpusGenerationPromptInput]):
+    """A basic prompt for corpus generation"""
+
+    system_prompt: str = (
+        "You are a provider of random factoids on topic requested by a user."
+        "Do not write a long essays, the response for given query should be a single sentence"
+        "For each query provide only a single fact about a given topic"
+        "Use as few tokens as possible"
+    )
+    user_prompt: str = "Provide factoids about {{ query }}"

ragbits/evaluate/dataset_generator/prompts/qa.py

@@ -0,0 +1,83 @@
+from pydantic import BaseModel
+
+from ragbits.core.prompt import Prompt
+
+
+class BasicAnswerGenInput(BaseModel):
+    """An input definition for basic answer generation task"""
+
+    chunk: str
+    question: str
+
+
+class BasicAnswerGenPrompt(Prompt[BasicAnswerGenInput, str]):
+    """A prompt clas for basic answers generation"""
+
+    system_prompt: str = (
+        "You are an AI assistant to answer the given question in the provided "
+        "evidence text. Do not mention any of these in the answer: 'in the "
+        "given text', 'in the provided information', etc. Users do not know "
+        "the passage source of the answer, so it should not be mentioned in "
+        "the answer. You can find the evidence from the given text about the "
+        "question, and you have to write a proper answer to the given question. "
+        "If you don't know the answer just say: I don't know."
+    )
+
+    user_prompt: str = "Text:\n<|text_start|>\n {{ chunk }} \n<|text_end|>\n\nQuestion:\n {{ question }} \n\nAnswer:"
+
+
+class PassagesGenInput(BaseModel):
+    """An input definition to passage generation prompt"""
+
+    question: str
+    basic_answer: str
+    chunk: str
+
+
+class PassagesGenPrompt(Prompt[PassagesGenInput, str]):
+    """A prompt class for passages generation"""
+
+    system_prompt: str = (
+        "You are an AI tasked with retrieving passages (one or many) from the "
+        "provided Chunk that contain information needed to generate the "
+        "provided Answer to the given Question.\n\nInstructions:\n1. Each "
+        "Passage MUST be VERBATIM and EXACT, without any modifications\n2. "
+        "Please provide the response in the form of a Python list. It should "
+        "begin with '[' and end with ']'\n3. You MUST start your answer with "
+        "'['\n4. The Chunk ALWAYS contains information needed to justify the "
+        "Answer\n5. Each passage must be as BRIEF as possible; DO NOT RETURN "
+        "FULL SENTENCES"
+    )
+
+    user_prompt: str = "Question:\n {{ question }} \nAnswer:\n {{ basic_answer }} \nChunk:\n {{ chunk }}\n\nPassages:"
+
+
+class QueryGenInput(BaseModel):
+    """An input definition for query generation prompt"""
+
+    chunk: str
+
+
+class QueryGenPrompt(Prompt[QueryGenInput, str]):
+    """A prompt class for query generation"""
+
+    system_prompt: str = (
+        "You're an AI tasked to convert Text into a factoid question. Factoid "
+        "questions are those seeking brief, factual information that can be "
+        "easily verified. They typically require a yes or no answer or a brief "
+        "explanation and often inquire about specific details such as dates, "
+        "names, places, or events.\n\nExamples of factoid questions include:\n"
+        "- What is the incoming shipment report?\n- What angle should I set my "
+        "ladder at?\n- What documents do I need to be a proof of transaction?\n\n"
+        "Instructions:\n1. Questions MUST BE extracted from given Text\n2. "
+        "Questions MUST BE as SHORT as possible\n3. Questions should be as "
+        "detailed as possible from Text\n4. Create questions that ask about "
+        "factual information from the Text\n5. Only return ONE question\n6. "
+        "Frame questions in a first-person, INFORMAL style, as if the employee "
+        "is seeking advice or clarification while working\n7. Do not mention any "
+        "of these in the questions: 'in the given text', 'in the provided "
+        "information', etc. Users do not know the passage source of the question, "
+        "so it should not be mentioned in the question."
+    )
+
+    user_prompt: str = "Text: {{ chunk }}\n\nGenerated Question from the Text:\n"

ragbits/evaluate/dataset_generator/tasks/corpus_generation.py

@@ -0,0 +1,67 @@
+import asyncio
+import sys
+from copy import deepcopy
+
+from distilabel.steps import StepInput, StepOutput
+from distilabel.steps.base import Step
+
+from ragbits.core.llms.base import LLM
+from ragbits.core.prompt import Prompt
+from ragbits.core.utils.config_handling import import_by_path
+
+module = sys.modules[__name__]
+
+
+class CorpusGenerationStep(Step):
+    """A step for corpus generation on given topics"""
+
+    def __init__(
+        self,
+        llm: LLM,
+        num_per_topic: int,
+        prompt_class: str | type[Prompt],
+    ):
+        super().__init__()
+        self._llm = llm
+        self._prompt_class = import_by_path(prompt_class, module) if isinstance(prompt_class, str) else prompt_class
+        self._num_per_topic = num_per_topic
+
+    @property
+    def inputs(self) -> list[str]:
+        """
+        A property defining input fields for a task
+        Returns:
+            list of input fields
+        """
+        return ["query"]
+
+    @property
+    def outputs(self) -> list[str]:
+        """
+        A property describing output fields for a step
+        Returns:
+            list of output fields
+        """
+        return ["chunk"]
+
+    def process(self, *inputs: StepInput) -> "StepOutput":
+        """
+        Generates the corpus data for a given topics
+        Args:
+            inputs: a topics on which the corpus data should be generated
+        Returns:
+            a generated corpus
+        """
+        result = asyncio.get_event_loop().run_until_complete(self._process_topics(topics=inputs[0]))
+        yield result
+
+    async def _process_topics(self, topics: list[dict]) -> list[dict]:
+        tasks = [self._process_topic(topic) for _ in range(self._num_per_topic) for topic in topics]
+        results = await asyncio.gather(*tasks)
+        return results
+
+    async def _process_topic(self, topic: dict) -> dict:
+        new_inp = deepcopy(topic)
+        prompt_inp = self._prompt_class.input_type(**{self.inputs[0]: new_inp[self.inputs[0]]})  # type: ignore
+        new_inp[self.outputs[0]] = await self._llm.generate(prompt=self._prompt_class(prompt_inp))
+        return new_inp

ragbits/evaluate/dataset_generator/tasks/filter/base.py

@@ -0,0 +1,43 @@
+from abc import ABC, abstractmethod
+
+from distilabel.steps import Step, StepInput, StepOutput
+
+from ..corpus_generation import CorpusGenerationStep
+from ..text_generation.base import BaseDistilabelTask
+
+
+class BaseFilter(Step, ABC):
+    """Base class for filtering the outputs of pipeline steps"""
+
+    def __init__(self, task: BaseDistilabelTask | CorpusGenerationStep):
+        super().__init__()
+        self._task = task
+
+    @property
+    def inputs(self) -> list[str]:
+        """
+        Property describing input fields for a filter
+        Returns:
+            list of input fields for a filter
+        """
+        return self._task.outputs
+
+    @property
+    def outputs(self) -> list[str]:
+        """
+        Property describing output fields for a filter
+        Returns:
+            list of output fields for a filter
+        """
+        return self._task.outputs
+
+    @abstractmethod
+    def process(self, *inputs: StepInput) -> "StepOutput":
+        """
+        Abstract method for filter step processing
+        Args:
+            inputs - inputs to a filter
+        Returns:
+            filtered outputs
+        """
+        pass

ragbits/evaluate/dataset_generator/tasks/filter/dont_know.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from distilabel.steps import StepInput, StepOutput
+
+from .base import BaseFilter
+
+DONT_KNOW_PHRASES: list[str] = [
+    "I don't know",
+    "I do not know",
+    "don't know",
+]
+
+
+class DontKnowFilter(BaseFilter):
+    """A class for basic rule-based filtering of don't know anwers"""
+
+    def process(self, *inputs: StepInput) -> "StepOutput":
+        """
+        Runs the basic rule-based filtering of the inputs
+        Args:
+            inputs - the outputs of some generation step
+        Returns:
+            outputs filtered to the ones that do not contain the pre-defined phrases
+        """
+        result = [
+            {input_type: input_[input_type] for input_type in input_}
+            for input_ in inputs[0]
+            if not self._is_dont_know(input_)
+        ]
+        yield result
+
+    @staticmethod
+    def _is_dont_know(input_: dict[str, Any]) -> bool:
+        return any(s.lower() in input_["basic_answer"].lower() for s in DONT_KNOW_PHRASES)

ragbits/evaluate/dataset_generator/tasks/text_generation/base.py

@@ -0,0 +1,66 @@
+import sys
+from abc import ABC, abstractmethod
+from typing import Any
+
+from distilabel.models import LLM
+from distilabel.steps.tasks import TextGeneration
+
+from ragbits.core.prompt import ChatFormat, Prompt
+from ragbits.core.utils.config_handling import import_by_path
+
+module = sys.modules[__name__]
+
+
+class BaseDistilabelTask(TextGeneration, ABC):
+    """Base class for distilabel TextGeneration tasks"""
+
+    def __init__(self, llm: LLM, inputs: list[str], outputs: list[str], prompt_class: str | type[Prompt]):
+        super().__init__(llm=llm)
+        self._inputs = inputs
+        self._outputs = outputs
+        self._prompt_class = import_by_path(prompt_class, module) if isinstance(prompt_class, str) else prompt_class
+
+    @property
+    def inputs(self) -> list[str]:
+        """
+        Property describing input fields for a task
+        Returns:
+            list of input fields for a task
+        """
+        return self._inputs
+
+    @property
+    def outputs(self) -> list[str]:
+        """
+        Property describing output fields of the task
+        Returns:
+            list of outputs for a task
+        """
+        return self._outputs
+
+    def format_input(self, input: dict[str, Any]) -> ChatFormat:
+        """
+        Formats the input data for generating a question based on the provided "chunk".
+
+        Args:
+            input: A dictionary containing a single "chunk" key with the text input.
+
+        Returns:
+            The formatted chat object containing the input for query generation.
+        """
+        chat = self._prompt_class(self._prompt_class.input_type(**input)).chat  # type: ignore
+        return chat
+
+    @abstractmethod
+    def format_output(self, output: str, input: dict[str, Any] | None = None) -> dict[str, str | list[str]]:
+        """
+        Formats the generated question into a structured dictionary with the original "chunk" input.
+
+        Args:
+            output: The generated question.
+            input: Optional; contains "chunk" key with the original input chunk.
+
+        Returns:
+            A dictionary containing "chunk" and "question".
+        """
+        pass

ragbits/evaluate/dataset_generator/tasks/text_generation/qa.py

@@ -0,0 +1,96 @@
+from typing import Any
+
+from distilabel.models import LLM
+
+from ragbits.evaluate.dataset_generator.tasks.text_generation.base import BaseDistilabelTask
+from ragbits.evaluate.dataset_generator.utils import get_closest_substring, get_passages_list
+
+
+class QueryGenTask(BaseDistilabelTask):
+    """
+    A task for generating a question based on a provided text chunk.
+    """
+
+    def __init__(self, llm: LLM, prompt_class: str):
+        super().__init__(llm=llm, inputs=["chunk"], outputs=["question", "chunk"], prompt_class=prompt_class)
+
+    def format_output(self, output: str, input: dict[str, Any] | None = None) -> dict[str, str | list[str]]:  # noqa: PLR6301
+        """
+        Formats the generated question into a structured dictionary with the original "chunk" input.
+
+        Args:
+            output: The generated question.
+            input: Optional; contains "chunk" key with the original input chunk.
+
+        Returns:
+            A dictionary containing "chunk" and "question".
+        """
+        return {"chunk": input["chunk"], "question": output}  # type: ignore
+
+
+class PassagesGenTask(BaseDistilabelTask):
+    """
+    A task for generating passages related to a specific question and answer from a text chunk.
+    """
+
+    should_get_matches: bool = False
+
+    def __init__(self, llm: LLM, prompt_class: str):
+        super().__init__(
+            llm=llm,
+            inputs=["chunk", "question", "basic_answer"],
+            outputs=["question", "chunk", "passages"],
+            prompt_class=prompt_class,
+        )
+
+    def format_output(self, output: str, input: dict[str, Any] | None = None) -> dict[str, str | list[str]]:
+        """
+        Formats the model's output into a structured dictionary with "question", "chunk", and "passages".
+        If `get_matches` is `True`, attempts to find the closest matches for each passage within the
+        provided chunk.
+
+        Args:
+            output: The raw output generated by the text generation model.
+            input: Required if `get_matches` is `True`, containing "chunk" and "question".
+
+        Returns:
+            A dictionary with "chunk", "question", and a list of "passages".
+        """
+        passages: list[str] = get_passages_list(output) or []
+
+        if self.should_get_matches:
+            matched_passages: list[str] = []
+
+            for passage in passages:
+                if passage in input["chunk"]:  # type: ignore
+                    matched_passages.append(passage)
+                else:
+                    matched_passage = get_closest_substring(input["chunk"], passage)  # type: ignore
+                    matched_passages.append(matched_passage)
+
+            return {"chunk": input["chunk"], "question": input["question"], "passages": matched_passages}  # type: ignore
+
+        return {"chunk": input["chunk"], "question": input["question"], "passages": passages}  # type: ignore
+
+
+class AnswerGenTask(BaseDistilabelTask):
+    """
+    A task for generating basic answers to questions based on a provided text chunk. This class extends
+    the `TextGeneration` task from the `distilabel` package.
+    """
+
+    def __init__(self, llm: LLM, prompt_class: str):
+        super().__init__(llm=llm, inputs=["chunk", "question"], outputs=["basic_answer"], prompt_class=prompt_class)
+
+    def format_output(self, output: str, input: dict[str, Any] | None = None) -> dict[str, str | list[str]]:  # noqa: PLR6301
+        """
+        Formats the model's output into a structured dictionary with the "basic_answer" key.
+
+        Args:
+            output: The raw output generated by the text generation model.
+            input: Optional; not typically used in this formatting.
+
+        Returns:
+            A dictionary with "basic_answer" as the key and the generated output as its value.
+        """
+        return {"basic_answer": output}

ragbits/evaluate/dataset_generator/utils.py

@@ -0,0 +1,43 @@
+import json
+import re
+import warnings
+from difflib import SequenceMatcher
+from itertools import combinations
+
+
+def get_closest_substring(long: str, short: str) -> str:
+    """
+    Finds the closest substring to short string in longer one
+    Args:
+        long: str - longer string
+        short: str - shorter string
+    Returns:
+        closest substring of longer
+    """
+    a, b = max(
+        combinations(re.finditer("|".join(short.split()), long), 2),
+        key=lambda c: SequenceMatcher(None, long[c[0].start() : c[1].end()], short).ratio(),
+    )
+    return long[a.start() : b.end()]
+
+
+def get_passages_list(raw_passages: str) -> list[str]:
+    """
+    Formats LLM output to list of passages
+    Args:
+        raw_passages: string representing raw passages returned by llm
+    Returns:
+        list of parsed passages
+    """
+    match = re.search(r"\[(.*?)\]", raw_passages, re.DOTALL)
+
+    if match:
+        passages_content = match.group(1)
+        try:
+            return json.loads("[" + passages_content + "]")
+        except (SyntaxError, ValueError):
+            warnings.warn("Unable to evaluate the passages content. Check the format.", category=UserWarning)
+            return []
+    else:
+        warnings.warn(message="No brackets found in the input string.", category=UserWarning)
+        return []