HuggingFace Model

rankify.generator.models.huggingface_model

BaseRAGModel

Bases: ABC

Base RAG Model for Retrieval-Augmented Generation (RAG).

This is an abstract base class for implementing LLM endpoints in rankify. It defines the interface for generating responses and optional embedding generation.

Methods:

Name	Description
generate	str, **kwargs) -> str: Abstract method to generate a response based on the given prompt.
embed	str, **kwargs) -> List[float]: Optional method to generate embeddings for the given text.

Notes

• This class serves as a blueprint for RAG models like OpenAIModel and HuggingFaceModel.
• The embed method is optional and can be implemented if needed.
• This class needs to be extended to include new LLM endpoints in Rankify.

Source code in rankify/generator/models/base_rag_model.py

class BaseRAGModel(ABC):
    """
    **Base RAG Model** for Retrieval-Augmented Generation (RAG).

    This is an abstract base class for implementing LLM endpoints in rankify. 
    It defines the interface for generating responses and optional embedding generation.

    Methods:
        generate(prompt: str, **kwargs) -> str:
            Abstract method to generate a response based on the given prompt.
        embed(text: str, **kwargs) -> List[float]:
            Optional method to generate embeddings for the given text.

    Notes:
        - This class serves as a blueprint for RAG models like `OpenAIModel` and `HuggingFaceModel`.
        - The `embed` method is optional and can be implemented if needed.
        - This class needs to be extended to include new LLM endpoints in Rankify.
    """

    @abstractmethod
    def generate(self, prompt: str, **kwargs) -> str:
        """Generate a response based on the given prompt."""
        pass

    def embed(self, text: str, **kwargs) -> List[float]:
        """Optional: Generate embeddings for the given text."""
        raise NotImplementedError("Embedding is not required for this implementation.")

generate(prompt, **kwargs) abstractmethod

Generate a response based on the given prompt.

Source code in rankify/generator/models/base_rag_model.py

@abstractmethod
def generate(self, prompt: str, **kwargs) -> str:
    """Generate a response based on the given prompt."""
    pass

embed(text, **kwargs)

Optional: Generate embeddings for the given text.

Source code in rankify/generator/models/base_rag_model.py

def embed(self, text: str, **kwargs) -> List[float]:
    """Optional: Generate embeddings for the given text."""
    raise NotImplementedError("Embedding is not required for this implementation.")

PromptGenerator

PromptGenerator for Retrieval-Augmented Generation (RAG).

This class manages prompt construction for different RAG methods and model types in Rankify. It selects and formats prompt templates based on the specified method and model, enabling flexible and consistent prompt generation.

Attributes:

Name	Type	Description
method	str	The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
prompt_template	PromptTemplate	The prompt template used for formatting prompts.
model_type	str	The type of model (e.g., "huggingface", "openai").

Methods:

Name	Description
generate_user_prompt	Generates a user prompt by formatting the question and contexts with the selected template.
_select_template	Selects the appropriate prompt template for the given method.

Notes

• Supports custom prompt templates via the custom_prompt argument.
• If no contexts are provided, generates prompts with only the question.
• Ensures consistent prompt formatting across different RAG methods and models.
• Automatically selects a default template if none is specified.

Source code in rankify/generator/prompt_generator.py

class PromptGenerator:
    """
    **PromptGenerator** for Retrieval-Augmented Generation (RAG).

    This class manages prompt construction for different RAG methods and model types in Rankify.
    It selects and formats prompt templates based on the specified method and model, enabling flexible and consistent prompt generation.

    Attributes:
        method (str): The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
        prompt_template (PromptTemplate): The prompt template used for formatting prompts.
        model_type (str): The type of model (e.g., "huggingface", "openai").

    Methods:
        generate_user_prompt(question, contexts, custom_prompt=None) -> str:
            Generates a user prompt by formatting the question and contexts with the selected template.
        _select_template(method) -> PromptTemplate:
            Selects the appropriate prompt template for the given method.

    Notes:
        - Supports custom prompt templates via the `custom_prompt` argument.
        - If no contexts are provided, generates prompts with only the question.
        - Ensures consistent prompt formatting across different RAG methods and models.
        - Automatically selects a default template if none is specified.
    """
    def __init__(self, method: str, model_type: str, prompt_template: Optional[PromptTemplate] = None):
        """
        Initialize the PromptGenerator.

        Args:
            method (str): The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
            model_type (str): The type of model (e.g., "huggingface", "openai").
            prompt_template (PromptTemplate, optional): Custom prompt template to use. If None, selects based on method.

        Notes:
            - If no custom template is provided, selects a default template for the specified method.
            - Stores the method, model type, and selected prompt template for prompt generation.
        """
        self.method = method
        if prompt_template is not None:
            self.prompt_template = prompt_template
        else:
            self.prompt_template = self._select_template(method)

    def _select_template(self, method: str) -> PromptTemplate:
        """
        Selects the appropriate prompt template for the given RAG method.

        Args:
            method (str): The RAG method or strategy.

        Returns:
            PromptTemplate: The selected prompt template.

        Notes:
            - If the method is not recognized, defaults to BASIC_RAG template.
        """
        method = method.lower()
        if method in PromptTemplate._value2member_map_:
            return PromptTemplate(method)
        return PromptTemplate.BASIC_RAG

    def generate_user_prompt(self, question: str, contexts: List[str], custom_prompt: Optional[str] = None) -> str:
        """
        Generates a user prompt by formatting the question and contexts with the selected template.

        Args:
            question (str): The question to be answered.
            contexts (List[str]): List of context passages to include in the prompt.
            custom_prompt (str, optional): Custom prompt template string. If provided, overrides the default template.

        Returns:
            str: The formatted prompt string.

        Notes:
            - If no contexts are provided, generates a prompt with only the question.
            - Custom prompt templates can use `{question}` and `{contexts}` placeholders.
            - Ensures consistent prompt formatting for all RAG methods and models.
        """
        if contexts is None:
            contexts = []
        context_str = "\n".join(contexts)
        if custom_prompt:
            return custom_prompt.format(question=question, contexts=context_str)
        template = DEFAULT_PROMPTS[self.prompt_template]
        return template.format(question=question, contexts=context_str)

__init__(method, model_type, prompt_template=None)

Initialize the PromptGenerator.

Parameters:

Name	Type	Description	Default
method	str	The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").	required
model_type	str	The type of model (e.g., "huggingface", "openai").	required
prompt_template	PromptTemplate	Custom prompt template to use. If None, selects based on method.	None

Notes

• If no custom template is provided, selects a default template for the specified method.
• Stores the method, model type, and selected prompt template for prompt generation.

Source code in rankify/generator/prompt_generator.py

def __init__(self, method: str, model_type: str, prompt_template: Optional[PromptTemplate] = None):
    """
    Initialize the PromptGenerator.

    Args:
        method (str): The RAG method or strategy (e.g., "basic-rag", "chain-of-thought-rag").
        model_type (str): The type of model (e.g., "huggingface", "openai").
        prompt_template (PromptTemplate, optional): Custom prompt template to use. If None, selects based on method.

    Notes:
        - If no custom template is provided, selects a default template for the specified method.
        - Stores the method, model type, and selected prompt template for prompt generation.
    """
    self.method = method
    if prompt_template is not None:
        self.prompt_template = prompt_template
    else:
        self.prompt_template = self._select_template(method)

generate_user_prompt(question, contexts, custom_prompt=None)

Generates a user prompt by formatting the question and contexts with the selected template.

Parameters:

Name	Type	Description	Default
question	str	The question to be answered.	required
contexts	List[str]	List of context passages to include in the prompt.	required
custom_prompt	str	Custom prompt template string. If provided, overrides the default template.	None

Returns:

Name	Type	Description
str	str	The formatted prompt string.

Notes

• If no contexts are provided, generates a prompt with only the question.
• Custom prompt templates can use {question} and {contexts} placeholders.
• Ensures consistent prompt formatting for all RAG methods and models.

Source code in rankify/generator/prompt_generator.py

def generate_user_prompt(self, question: str, contexts: List[str], custom_prompt: Optional[str] = None) -> str:
    """
    Generates a user prompt by formatting the question and contexts with the selected template.

    Args:
        question (str): The question to be answered.
        contexts (List[str]): List of context passages to include in the prompt.
        custom_prompt (str, optional): Custom prompt template string. If provided, overrides the default template.

    Returns:
        str: The formatted prompt string.

    Notes:
        - If no contexts are provided, generates a prompt with only the question.
        - Custom prompt templates can use `{question}` and `{contexts}` placeholders.
        - Ensures consistent prompt formatting for all RAG methods and models.
    """
    if contexts is None:
        contexts = []
    context_str = "\n".join(contexts)
    if custom_prompt:
        return custom_prompt.format(question=question, contexts=context_str)
    template = DEFAULT_PROMPTS[self.prompt_template]
    return template.format(question=question, contexts=context_str)

HuggingFaceModel

Bases: BaseRAGModel

Hugging Face Model for Retrieval-Augmented Generation (RAG).

This class integrates Hugging Face's pretrained models for text generation in a RAG pipeline. It uses the Hugging Face Transformers library for tokenization and model inference.

Attributes:

Name	Type	Description
model_name	str	Name of the Hugging Face model.
tokenizer		Tokenizer instance for encoding input text.
model		Pretrained Hugging Face model for text generation.
prompt_generator	PromptGenerator	Instance for generating prompts.
stop_at_period	bool	If True, cuts generated answer at the first period.

Notes

• This model uses Hugging Face's Transformers library for text generation.
• Default generation parameters like max_length and temperature can be overridden.

Source code in rankify/generator/models/huggingface_model.py

class HuggingFaceModel(BaseRAGModel):
    """
    **Hugging Face Model** for Retrieval-Augmented Generation (RAG).

    This class integrates Hugging Face's pretrained models for text generation in a RAG pipeline. 
    It uses the Hugging Face Transformers library for tokenization and model inference.

    Attributes:
        model_name (str): Name of the Hugging Face model.
        tokenizer: Tokenizer instance for encoding input text.
        model: Pretrained Hugging Face model for text generation.
        prompt_generator (PromptGenerator): Instance for generating prompts.
        stop_at_period (bool): If True, cuts generated answer at the first period.

    Notes:
        - This model uses Hugging Face's Transformers library for text generation.
        - Default generation parameters like `max_length` and `temperature` can be overridden.
    """

    def __init__(self, model_name: str, tokenizer, model, prompt_generator: PromptGenerator, stop_at_period: bool = False):
        self.model_name = model_name
        self.tokenizer = tokenizer
        self.model = model
        self.prompt_generator = prompt_generator
        self.stop_at_period = stop_at_period

    def generate(self, prompt: str, **kwargs):
        """
        Generates a response using the Hugging Face model and returns the answer(s).

        Args:
            prompt (str): The input prompt for generation.
            **kwargs: Optional generation parameters, such as:
                - max_new_tokens (int): Maximum number of new tokens to generate (default: 64).
                - do_sample (bool): Whether to use sampling (default: True).
                - num_return_sequences (int): Number of answers to generate (default: 1).
                - eos_token_id (int): End-of-sequence token ID (default: tokenizer.eos_token_id).
                - pad_token_id (int): Padding token ID (default: tokenizer.eos_token_id).
                - temperature (float): Sampling temperature (default: 0.1).
                - top_p (float): Nucleus sampling parameter (default: 1.0).

        Returns:
            str or List[str]: The generated answer(s). If `num_return_sequences` > 1, returns a list of answers.

        Notes:
            - The answer is post-processed to remove the prompt and extra whitespace.
            - If `stop_at_period` is True, the answer is truncated at the first period.
            - All generation parameters can be overridden via `kwargs`.

        Example:
            ```python
            answer = model.generate("What is the capital of France?", max_new_tokens=32)
            ```
        """
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)

        kwargs.setdefault("max_new_tokens", 64)
        kwargs.setdefault("do_sample", True)
        kwargs.setdefault("num_return_sequences", 1)
        kwargs.setdefault("eos_token_id", self.tokenizer.eos_token_id)
        kwargs.setdefault("pad_token_id", self.tokenizer.eos_token_id)
        kwargs.setdefault("temperature", 0.1)
        kwargs.setdefault("top_p", 1.0)

        outputs = self.model.generate(**inputs, **kwargs)

        def clean_answer(text):
            answer = text[len(prompt):].strip()
            answer = answer.split("\n")[0].strip()
            if self.stop_at_period:
                idx = answer.find(".")
                if idx != -1:
                    return answer[:idx+1].strip()
            return answer

        if kwargs.get("num_return_sequences", 1) > 1:
            return [clean_answer(self.tokenizer.decode(output, skip_special_tokens=True)) for output in outputs]
        else:
            return clean_answer(self.tokenizer.decode(outputs[0], skip_special_tokens=True))

generate(prompt, **kwargs)

Generates a response using the Hugging Face model and returns the answer(s).

Parameters:

Name	Type	Description	Default
prompt	str	The input prompt for generation.	required
**kwargs		Optional generation parameters, such as: - max_new_tokens (int): Maximum number of new tokens to generate (default: 64). - do_sample (bool): Whether to use sampling (default: True). - num_return_sequences (int): Number of answers to generate (default: 1). - eos_token_id (int): End-of-sequence token ID (default: tokenizer.eos_token_id). - pad_token_id (int): Padding token ID (default: tokenizer.eos_token_id). - temperature (float): Sampling temperature (default: 0.1). - top_p (float): Nucleus sampling parameter (default: 1.0).	{}

Returns:

Type	Description
	str or List[str]: The generated answer(s). If num_return_sequences > 1, returns a list of answers.

Notes

• The answer is post-processed to remove the prompt and extra whitespace.
• If stop_at_period is True, the answer is truncated at the first period.
• All generation parameters can be overridden via kwargs.

Example

answer = model.generate("What is the capital of France?", max_new_tokens=32)

Source code in rankify/generator/models/huggingface_model.py

def generate(self, prompt: str, **kwargs):
    """
    Generates a response using the Hugging Face model and returns the answer(s).

    Args:
        prompt (str): The input prompt for generation.
        **kwargs: Optional generation parameters, such as:
            - max_new_tokens (int): Maximum number of new tokens to generate (default: 64).
            - do_sample (bool): Whether to use sampling (default: True).
            - num_return_sequences (int): Number of answers to generate (default: 1).
            - eos_token_id (int): End-of-sequence token ID (default: tokenizer.eos_token_id).
            - pad_token_id (int): Padding token ID (default: tokenizer.eos_token_id).
            - temperature (float): Sampling temperature (default: 0.1).
            - top_p (float): Nucleus sampling parameter (default: 1.0).

    Returns:
        str or List[str]: The generated answer(s). If `num_return_sequences` > 1, returns a list of answers.

    Notes:
        - The answer is post-processed to remove the prompt and extra whitespace.
        - If `stop_at_period` is True, the answer is truncated at the first period.
        - All generation parameters can be overridden via `kwargs`.

    Example:
        ```python
        answer = model.generate("What is the capital of France?", max_new_tokens=32)
        ```
    """
    inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)

    kwargs.setdefault("max_new_tokens", 64)
    kwargs.setdefault("do_sample", True)
    kwargs.setdefault("num_return_sequences", 1)
    kwargs.setdefault("eos_token_id", self.tokenizer.eos_token_id)
    kwargs.setdefault("pad_token_id", self.tokenizer.eos_token_id)
    kwargs.setdefault("temperature", 0.1)
    kwargs.setdefault("top_p", 1.0)

    outputs = self.model.generate(**inputs, **kwargs)

    def clean_answer(text):
        answer = text[len(prompt):].strip()
        answer = answer.split("\n")[0].strip()
        if self.stop_at_period:
            idx = answer.find(".")
            if idx != -1:
                return answer[:idx+1].strip()
        return answer

    if kwargs.get("num_return_sequences", 1) > 1:
        return [clean_answer(self.tokenizer.decode(output, skip_special_tokens=True)) for output in outputs]
    else:
        return clean_answer(self.tokenizer.decode(outputs[0], skip_special_tokens=True))