vLLM Model

rankify.generator.models.vllm_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)

VLLMModel

Bases: BaseRAGModel

vLLM Model for Retrieval-Augmented Generation (RAG).

This class integrates vLLM's API for text generation in a RAG pipeline. It uses the vLLM library to generate responses based on input prompts.

Attributes:

Name	Type	Description
model_name	str	Name of the vLLM model.
prompt_generator	PromptGenerator	Instance for generating prompts.
client	LLM	Client for interacting with the vLLM library.

Notes

• This model uses vLLM's library for text generation.
• It supports additional parameters like max_tokens and temperature.
• Specialty: vLLM requires a sampling_params dictionary to control decoding and sampling behavior, allowing fine-grained control over generation (e.g., max_tokens, temperature, top_p, etc.).

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

class VLLMModel(BaseRAGModel):
    """
    **vLLM Model** for Retrieval-Augmented Generation (RAG).

    This class integrates vLLM's API for text generation in a RAG pipeline. 
    It uses the vLLM library to generate responses based on input prompts.

    Attributes:
        model_name (str): Name of the vLLM model.
        prompt_generator (PromptGenerator): Instance for generating prompts.
        client (LLM): Client for interacting with the vLLM library.

    Notes:
        - This model uses vLLM's library for text generation.
        - It supports additional parameters like `max_tokens` and `temperature`.
        - **Specialty:** vLLM requires a `sampling_params` dictionary to control decoding and sampling behavior, 
          allowing fine-grained control over generation (e.g., `max_tokens`, `temperature`, `top_p`, etc.).
    """
    def __init__(self, model_name: str, prompt_generator: PromptGenerator, **kwargs):
        """
        Initialize the VLLMModel with the vLLM client.

        :param model_name: Name of the vLLM model.
        :param prompt_generator: Instance of PromptGenerator for generating prompts.
        :param device: Device to run the model on (default: "cuda").
        """
        self.model_name = model_name
        self.prompt_generator = prompt_generator
        self.llm = LLM(model=model_name, **kwargs)

    def generate(self, prompt: str, **kwargs) -> str:
        """
        Generate a response using vLLM's API.

        Args:
            prompt (str): The input prompt for the model.
            sampling_params (dict, required in kwargs): Dictionary of sampling parameters for vLLM generation, such as:
                - max_tokens (int): Maximum number of tokens to generate.
                - temperature (float): Sampling temperature.
                - top_p (float): Nucleus sampling parameter.
                - other vLLM-supported generation options.
            **kwargs: Additional parameters for the vLLM API call.

        Returns:
            str: The generated response as a string.

        Notes:
            - `sampling_params` controls the decoding and sampling behavior of the vLLM model.
            - Typical keys include `max_tokens`, `temperature`, `top_p`, etc.
            - See vLLM documentation for all supported sampling parameters.

        Example:
            ```python
            answer = model.generate(
                "What is the capital of France?",
                sampling_params={"max_tokens": 64, "temperature": 0.7, "top_p": 0.9}
            )
        ```
        """
        # Call the vLLM API using the LLM client
        response = self.llm.generate(prompt, kwargs["sampling_params"])

        if response and response[0].outputs:
            # Extract the generated text from the response
            generated_text = response[0].outputs[0].text.strip()
        else:
            generated_text = ""

        return generated_text

__init__(model_name, prompt_generator, **kwargs)

Initialize the VLLMModel with the vLLM client.

:param model_name: Name of the vLLM model. :param prompt_generator: Instance of PromptGenerator for generating prompts. :param device: Device to run the model on (default: "cuda").

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

def __init__(self, model_name: str, prompt_generator: PromptGenerator, **kwargs):
    """
    Initialize the VLLMModel with the vLLM client.

    :param model_name: Name of the vLLM model.
    :param prompt_generator: Instance of PromptGenerator for generating prompts.
    :param device: Device to run the model on (default: "cuda").
    """
    self.model_name = model_name
    self.prompt_generator = prompt_generator
    self.llm = LLM(model=model_name, **kwargs)

generate(prompt, **kwargs)

Generate a response using vLLM's API.

Parameters:

Name	Type	Description	Default
prompt	str	The input prompt for the model.	required
sampling_params	(dict, required in kwargs)	Dictionary of sampling parameters for vLLM generation, such as: - max_tokens (int): Maximum number of tokens to generate. - temperature (float): Sampling temperature. - top_p (float): Nucleus sampling parameter. - other vLLM-supported generation options.	required
**kwargs		Additional parameters for the vLLM API call.	{}

Returns:

Name	Type	Description
str	str	The generated response as a string.

Notes

• sampling_params controls the decoding and sampling behavior of the vLLM model.
• Typical keys include max_tokens, temperature, top_p, etc.
• See vLLM documentation for all supported sampling parameters.

Example

```python answer = model.generate( "What is the capital of France?", sampling_params={"max_tokens": 64, "temperature": 0.7, "top_p": 0.9} )

```

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

def generate(self, prompt: str, **kwargs) -> str:
    """
    Generate a response using vLLM's API.

    Args:
        prompt (str): The input prompt for the model.
        sampling_params (dict, required in kwargs): Dictionary of sampling parameters for vLLM generation, such as:
            - max_tokens (int): Maximum number of tokens to generate.
            - temperature (float): Sampling temperature.
            - top_p (float): Nucleus sampling parameter.
            - other vLLM-supported generation options.
        **kwargs: Additional parameters for the vLLM API call.

    Returns:
        str: The generated response as a string.

    Notes:
        - `sampling_params` controls the decoding and sampling behavior of the vLLM model.
        - Typical keys include `max_tokens`, `temperature`, `top_p`, etc.
        - See vLLM documentation for all supported sampling parameters.

    Example:
        ```python
        answer = model.generate(
            "What is the capital of France?",
            sampling_params={"max_tokens": 64, "temperature": 0.7, "top_p": 0.9}
        )
    ```
    """
    # Call the vLLM API using the LLM client
    response = self.llm.generate(prompt, kwargs["sampling_params"])

    if response and response[0].outputs:
        # Extract the generated text from the response
        generated_text = response[0].outputs[0].text.strip()
    else:
        generated_text = ""

    return generated_text