Skip to content

Model Factory

rankify.generator.models.model_factory

FiDModel

Bases: BaseRAGModel

FiD (Fusion-in-Decoder) Generator for Open-Domain Question Answering.

Fusion-in-Decoder (FiD) is a retrieval-augmented generation (RAG) model that aggregates information from multiple retrieved passages to generate context-aware answers. Inside Rankifys architecture, this model is an exception, it subclasses BaseRAGModel but also includes the logic of the FiD RAG method, since this RAG technique relies on the full transformer architecture.

Attributes:

Name Type Description
device str

Device used for inference ('cuda' or 'cpu').
model_path str

Path to the downloaded and extracted model.
tokenizer T5Tokenizer

T5 tokenizer for text encoding.
model FiDT5

Pretrained FiD model for text generation.
max_length int

Maximum length of the generated output (default: 50 tokens).
References
  • Izacard & Grave Leveraging Passage Retrieval with Generative Models for Open-Domain QA
    Paper
See Also
  • BaseRAGModel: Parent class for FiDModel.
  • RAG-based QA Models: FiD falls under retrieval-augmented generation.
Example 
from rankify.dataset.dataset import Document, Question, Answer, Context
from rankify.generator.generator import Generator

# Define question and answer
question = Question("What is the capital of France?")
answers = Answer([""])
contexts = [
    Context(id=1, title="France", text="The capital of France is Paris.", score=0.9),
    Context(id=2, title="Germany", text="Berlin is the capital of Germany.", score=0.5)
]

# Construct document
doc = Document(question=question, answers=answers, contexts=contexts)

# Initialize Generator (e.g., Meta Llama)
generator = Generator(method="fid", model_name='nq_reader_base', backend="fid")

# Generate answer
generated_answers = generator.generate([doc])
print(generated_answers)  # Output: ["Paris"]
Notes
  • FiD combines multiple passages to generate better responses.
  • It integrates seamlessly with the Generator class.
  • Uses retrieval-augmented generation (RAG) techniques for QA.
Source code in rankify/generator/models/fid_model.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
class FiDModel(BaseRAGModel):
    """
    **FiD (Fusion-in-Decoder) Generator** for Open-Domain Question Answering.


    **Fusion-in-Decoder (FiD)** is a **retrieval-augmented generation (RAG) model** that 
    aggregates information from **multiple retrieved passages** to generate context-aware answers.
    Inside Rankifys architecture, this model is an exception, it subclasses `BaseRAGModel` but also includes the logic
    of the FiD RAG method, since this RAG technique relies on the full transformer architecture.

    Attributes:
        device (str): Device used for inference (`'cuda'` or `'cpu'`).
        model_path (str): Path to the downloaded and extracted model.
        tokenizer (transformers.T5Tokenizer): T5 tokenizer for text encoding.
        model (FiDT5): Pretrained FiD model for text generation.
        max_length (int): Maximum length of the generated output (default: 50 tokens).

    References:
        - **Izacard & Grave** *Leveraging Passage Retrieval with Generative Models for Open-Domain QA*  
          [Paper](https://arxiv.org/abs/2007.01282)

    See Also:
        - `BaseRAGModel`: Parent class for FiDModel.
        - `RAG-based QA Models`: FiD falls under retrieval-augmented generation.

    Example:
        ```python
        from rankify.dataset.dataset import Document, Question, Answer, Context
        from rankify.generator.generator import Generator

        # Define question and answer
        question = Question("What is the capital of France?")
        answers = Answer([""])
        contexts = [
            Context(id=1, title="France", text="The capital of France is Paris.", score=0.9),
            Context(id=2, title="Germany", text="Berlin is the capital of Germany.", score=0.5)
        ]

        # Construct document
        doc = Document(question=question, answers=answers, contexts=contexts)

        # Initialize Generator (e.g., Meta Llama)
        generator = Generator(method="fid", model_name='nq_reader_base', backend="fid")

        # Generate answer
        generated_answers = generator.generate([doc])
        print(generated_answers)  # Output: ["Paris"]
        ```

    Notes:
        - FiD **combines multiple passages** to generate better responses.
        - It integrates **seamlessly** with the `Generator` class.
        - Uses **retrieval-augmented generation (RAG)** techniques for QA.
    """


    MODEL_URLS = {
        "nq_reader_base": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/nq_reader_base.tar.gz",
        "nq_reader_large": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/nq_reader_large.tar.gz",
        "tqa_reader_base": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/tqa_reader_base.tar.gz",
        "tqa_reader_large": "https://dl.fbaipublicfiles.com/FiD/pretrained_models/tqa_reader_large.tar.gz"
    }

    CACHE_DIR = os.environ.get("RERANKING_CACHE_DIR", "./cache")

    def __init__(self, method: str = "fid", model_name: str = "nq_reader_base", **kwargs):
        """
        Initializes the FiDModel for retrieval-augmented generation.

        Args:
            method (str): The generator type (default: `"fid"`).
            model_name (str): The specific FiD model to use (e.g., `"nq_reader_base"`).
            max_length (int): Maximum length of generated answers (default: 50).
            device (str, optional): Device for inference (`'cuda'` or `'cpu'`). If None, auto-detect.
            **kwargs: Additional parameters for model configuration.

        Sets up device, downloads and loads the model and tokenizer, and configures generation parameters.

        Example:
            ```python
            generator = FiDModel(method="fid", model_name="nq_reader_base", max_length=64, device="cuda")
            ```
        """
        self.method = method
        self.model_name = model_name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model_path = self._download_model()
        self.tokenizer, self.model = self._load_model()
        self.max_length = kwargs.get("max_length", 50)  # Default max length for answers

    def _download_model(self) -> str:
        """
        Downloads and extracts the FiD model if not already present.

        Returns:
            str: The directory path where the model is stored.

        Raises:
            ValueError: If an unknown FiD model is specified.

        Example:
            ```python
            model_path = generator._download_model()
            ```
        """
        model_url = self.MODEL_URLS.get(self.model_name)
        if not model_url:
            raise ValueError(f"Unknown FiD model: {self.model_name}")

        model_dir = os.path.join(self.CACHE_DIR, "generator", "fid", self.model_name)
        if not os.path.exists(model_dir):
            ModelDownloader.download_and_extract(model_url, model_dir)

        return model_dir

    def _load_model(self):
        """
        Loads the FiD model and tokenizer.

        Returns:
            tuple: A tuple containing the tokenizer and the FiD model.

        Example:
            ```python
            tokenizer, model = generator._load_model()
            ```
        """
        tokenizer = transformers.T5Tokenizer.from_pretrained('t5-base', return_dict=False)
        model = FiDT5.from_pretrained(self.model_path)
        model = model.to(self.device)
        model.eval()
        return tokenizer, model

    def _prepare_dataloader(self, documents: list[Document]):
        """
        Converts `Document` objects into a FiD-compatible dataset.

        Args:
            documents (List[Document]): A list of documents with queries and retrieved contexts.

        Returns:
            tuple: A tuple containing a DataLoader and Dataset.

        Example:
            ```python
            dataloader, dataset = generator._prepare_dataloader(documents)
            ```
        """
        examples = []
        for doc in documents:
            example = {
                "question": doc.question.question,
                "answers": doc.answers.answers,
                "ctxs": [{"title": ctx.title, "text": ctx.text, "score": ctx.score} for ctx in doc.contexts]
            }
            examples.append(example)

        eval_dataset = Dataset(examples, n_context=len(documents[0].contexts))
        eval_sampler = SequentialSampler(eval_dataset)
        collator_function = Collator(512, self.tokenizer)  # 512 is max text length

        eval_dataloader = DataLoader(
            eval_dataset,
            sampler=eval_sampler,
            batch_size=4,  # Modify as needed
            num_workers=4,
            collate_fn=collator_function
        )

        return eval_dataloader, eval_dataset

    def generate(self, documents: list[Document]) -> list[str]:
        """
        Generates answers for a list of documents.

        Args:
            documents (List[Document]): A list of documents with queries and retrieved contexts.

        Returns:
            List[str]: A list of generated answers.

        Example:
            ```python
            generator = FiDGenerator(method="fid", model_name="nq_reader_base")
            documents = [Document(question=Question("Who wrote Hamlet?"))]
            print(generator.generate(documents))  # ['William Shakespeare wrote Hamlet in the early 1600s.']
            ```
        """
        eval_dataloader, eval_dataset = self._prepare_dataloader(documents)

        results = []
        with torch.no_grad():
            for batch in eval_dataloader:
                (idx, _, _, context_ids, context_mask) = batch

                outputs = self.model.generate(
                    input_ids=context_ids.to(self.device),
                    attention_mask=context_mask.to(self.device),
                    max_length=self.max_length
                )

                for k, output in enumerate(outputs):
                    answer = self.tokenizer.decode(output, skip_special_tokens=True)
                    results.append(answer)

        return results

__init__(method='fid', model_name='nq_reader_base', **kwargs)

Initializes the FiDModel for retrieval-augmented generation.

Parameters:

Name Type Description Default
method str

The generator type (default: "fid").

 'fid'
model_name str

The specific FiD model to use (e.g., "nq_reader_base").

 'nq_reader_base'
max_length int

Maximum length of generated answers (default: 50).

 required
device str

Device for inference ('cuda' or 'cpu'). If None, auto-detect.

 required
**kwargs

Additional parameters for model configuration.

 {}

Sets up device, downloads and loads the model and tokenizer, and configures generation parameters.

Example 
generator = FiDModel(method="fid", model_name="nq_reader_base", max_length=64, device="cuda")
Source code in rankify/generator/models/fid_model.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def __init__(self, method: str = "fid", model_name: str = "nq_reader_base", **kwargs):
    """
    Initializes the FiDModel for retrieval-augmented generation.

    Args:
        method (str): The generator type (default: `"fid"`).
        model_name (str): The specific FiD model to use (e.g., `"nq_reader_base"`).
        max_length (int): Maximum length of generated answers (default: 50).
        device (str, optional): Device for inference (`'cuda'` or `'cpu'`). If None, auto-detect.
        **kwargs: Additional parameters for model configuration.

    Sets up device, downloads and loads the model and tokenizer, and configures generation parameters.

    Example:
        ```python
        generator = FiDModel(method="fid", model_name="nq_reader_base", max_length=64, device="cuda")
        ```
    """
    self.method = method
    self.model_name = model_name
    self.device = "cuda" if torch.cuda.is_available() else "cpu"
    self.model_path = self._download_model()
    self.tokenizer, self.model = self._load_model()
    self.max_length = kwargs.get("max_length", 50)  # Default max length for answers

generate(documents)

Generates answers for a list of documents.

Parameters:

Name Type Description Default
documents List[Document]

A list of documents with queries and retrieved contexts.

 required

Returns:

Type Description
list[str]

List[str]: A list of generated answers.
Example 
generator = FiDGenerator(method="fid", model_name="nq_reader_base")
documents = [Document(question=Question("Who wrote Hamlet?"))]
print(generator.generate(documents))  # ['William Shakespeare wrote Hamlet in the early 1600s.']
Source code in rankify/generator/models/fid_model.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def generate(self, documents: list[Document]) -> list[str]:
    """
    Generates answers for a list of documents.

    Args:
        documents (List[Document]): A list of documents with queries and retrieved contexts.

    Returns:
        List[str]: A list of generated answers.

    Example:
        ```python
        generator = FiDGenerator(method="fid", model_name="nq_reader_base")
        documents = [Document(question=Question("Who wrote Hamlet?"))]
        print(generator.generate(documents))  # ['William Shakespeare wrote Hamlet in the early 1600s.']
        ```
    """
    eval_dataloader, eval_dataset = self._prepare_dataloader(documents)

    results = []
    with torch.no_grad():
        for batch in eval_dataloader:
            (idx, _, _, context_ids, context_mask) = batch

            outputs = self.model.generate(
                input_ids=context_ids.to(self.device),
                attention_mask=context_mask.to(self.device),
                max_length=self.max_length
            )

            for k, output in enumerate(outputs):
                answer = self.tokenizer.decode(output, skip_special_tokens=True)
                results.append(answer)

    return results

LitellmModel

Bases: BaseRAGModel

LiteLLM Model for Retrieval-Augmented Generation (RAG).

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

Attributes:

Name Type Description
model_name str

Name of the LiteLLM model.
prompt_generator PromptGenerator

Instance for generating prompts.
client LitellmClient

Client for interacting with the LiteLLM API.
Notes
  • This model uses LiteLLM's API for text generation.
  • It supports additional parameters like max_tokens and temperature.
Source code in rankify/generator/models/litellm_model.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class LitellmModel(BaseRAGModel):
    """
    **LiteLLM Model** for Retrieval-Augmented Generation (RAG).

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

    Attributes:
        model_name (str): Name of the LiteLLM model.
        prompt_generator (PromptGenerator): Instance for generating prompts.
        client (LitellmClient): Client for interacting with the LiteLLM API.

    Notes:
        - This model uses LiteLLM's API for text generation.
        - It supports additional parameters like `max_tokens` and `temperature`.
    """
    def __init__(self, model_name: str, api_key: str, prompt_generator: PromptGenerator):
        """
        Initialize the LitellmModel with the LitellmClient.

        :param model_name: Name of the LiteLLM model.
        :param api_key: API key for LiteLLM.
        :param prompt_generator: Instance of PromptGenerator for generating prompts.
        """
        self.model_name = model_name
        self.prompt_generator = prompt_generator

        self.client = LitellmClient(keys=api_key)

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

        Args:
            prompt (str): The input prompt for the model.
            **kwargs: Additional parameters for the LiteLLM API call, such as:
                - model (str): Model name to use (default: self.model_name).
                - max_tokens (int): Maximum number of tokens to generate (default: 128).
                - temperature (float): Sampling temperature (default: 0.7).

        Returns:
            str: The generated response as a string.

        Notes:
            - Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
            - All generation parameters can be overridden via `kwargs`.

        Example:
            ```python
            answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
            ```
        """

        # Set default parameters for the LiteLLM API call
        kwargs.setdefault("model", self.model_name)
        kwargs.setdefault("max_tokens", 128)
        kwargs.setdefault("temperature", 0.7)

        # Call the LiteLLM API using the LitellmClient
        response = self.client.chat(messages=[{"role": "user", "content": prompt}], return_text=True, **kwargs)
        return response

__init__(model_name, api_key, prompt_generator)

Initialize the LitellmModel with the LitellmClient.

:param model_name: Name of the LiteLLM model. :param api_key: API key for LiteLLM. :param prompt_generator: Instance of PromptGenerator for generating prompts.

Source code in rankify/generator/models/litellm_model.py 
21
22
23
24
25
26
27
28
29
30
31
32
def __init__(self, model_name: str, api_key: str, prompt_generator: PromptGenerator):
    """
    Initialize the LitellmModel with the LitellmClient.

    :param model_name: Name of the LiteLLM model.
    :param api_key: API key for LiteLLM.
    :param prompt_generator: Instance of PromptGenerator for generating prompts.
    """
    self.model_name = model_name
    self.prompt_generator = prompt_generator

    self.client = LitellmClient(keys=api_key)

generate(prompt, **kwargs)

Generate a response using LiteLLM's API.

Parameters:

Name Type Description Default
prompt str

The input prompt for the model.

 required
**kwargs

Additional parameters for the LiteLLM API call, such as: - model (str): Model name to use (default: self.model_name). - max_tokens (int): Maximum number of tokens to generate (default: 128). - temperature (float): Sampling temperature (default: 0.7).

 {}

Returns:

Name Type Description
str str

The generated response as a string.
Notes
  • Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
  • All generation parameters can be overridden via kwargs.
Example 
answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
Source code in rankify/generator/models/litellm_model.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def generate(self, prompt: str, **kwargs) -> str:
    """
    Generate a response using LiteLLM's API.

    Args:
        prompt (str): The input prompt for the model.
        **kwargs: Additional parameters for the LiteLLM API call, such as:
            - model (str): Model name to use (default: self.model_name).
            - max_tokens (int): Maximum number of tokens to generate (default: 128).
            - temperature (float): Sampling temperature (default: 0.7).

    Returns:
        str: The generated response as a string.

    Notes:
        - Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
        - All generation parameters can be overridden via `kwargs`.

    Example:
        ```python
        answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
        ```
    """

    # Set default parameters for the LiteLLM API call
    kwargs.setdefault("model", self.model_name)
    kwargs.setdefault("max_tokens", 128)
    kwargs.setdefault("temperature", 0.7)

    # Call the LiteLLM API using the LitellmClient
    response = self.client.chat(messages=[{"role": "user", "content": prompt}], return_text=True, **kwargs)
    return response

OpenAIModel

Bases: BaseRAGModel

OpenAI Model for Retrieval-Augmented Generation (RAG).

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

Attributes:

Name Type Description
model_name str

Name of the OpenAI model (e.g., "gpt-3.5-turbo").
prompt_generator PromptGenerator

Instance for generating prompts.
client OpenaiClient

Client for interacting with the OpenAI API.
Notes
  • This model uses OpenAI's GPT models for text generation.
  • It supports additional parameters like max_tokens and temperature.
Source code in rankify/generator/models/openai_model.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class OpenAIModel(BaseRAGModel):
    """
    **OpenAI Model** for Retrieval-Augmented Generation (RAG).

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

    Attributes:
        model_name (str): Name of the OpenAI model (e.g., "gpt-3.5-turbo").
        prompt_generator (PromptGenerator): Instance for generating prompts.
        client (OpenaiClient): Client for interacting with the OpenAI API.

    Notes:
        - This model uses OpenAI's GPT models for text generation.
        - It supports additional parameters like `max_tokens` and `temperature`.
    """
    def __init__(self, model_name: str, api_keys: list, prompt_generator: PromptGenerator, base_url: str = None):
        """
        Initialize the OpenAIModel with the OpenaiClient.

        :param model_name: Name of the OpenAI model (e.g., "gpt-3.5-turbo").
        :param api_keys: List of API keys for OpenAI.
        :param prompt_generator: Instance of PromptGenerator for generating prompts.
        :param base_url: Optional base URL for the OpenAI API.
        """
        self.model_name = model_name
        self.prompt_generator = prompt_generator

        self.client = OpenaiClient(keys=api_keys, base_url=base_url)

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

        Args:
            prompt (str): The input prompt for the model.
            **kwargs: Additional parameters for the OpenAI API call, such as:
                - model (str): Model name to use (default: self.model_name).
                - max_tokens (int): Maximum number of tokens to generate (default: 128).
                - temperature (float): Sampling temperature (default: 0.7).

        Returns:
            str: The generated response as a string.

        Notes:
            - Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
            - All generation parameters can be overridden via `kwargs`.

        Example:
            ```python
            answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
            ```
        """
        # Set default parameters for the OpenAI API call
        kwargs.setdefault("model", self.model_name)
        kwargs.setdefault("max_tokens", 128)
        kwargs.setdefault("temperature", 0.7)

        # Call the OpenAI API using the OpenaiClient
        response = self.client.chat(messages=[{"role": "user", "content": prompt}], return_text=True, **kwargs)
        return response

__init__(model_name, api_keys, prompt_generator, base_url=None)

Initialize the OpenAIModel with the OpenaiClient.

:param model_name: Name of the OpenAI model (e.g., "gpt-3.5-turbo"). :param api_keys: List of API keys for OpenAI. :param prompt_generator: Instance of PromptGenerator for generating prompts. :param base_url: Optional base URL for the OpenAI API.

Source code in rankify/generator/models/openai_model.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
def __init__(self, model_name: str, api_keys: list, prompt_generator: PromptGenerator, base_url: str = None):
    """
    Initialize the OpenAIModel with the OpenaiClient.

    :param model_name: Name of the OpenAI model (e.g., "gpt-3.5-turbo").
    :param api_keys: List of API keys for OpenAI.
    :param prompt_generator: Instance of PromptGenerator for generating prompts.
    :param base_url: Optional base URL for the OpenAI API.
    """
    self.model_name = model_name
    self.prompt_generator = prompt_generator

    self.client = OpenaiClient(keys=api_keys, base_url=base_url)

generate(prompt, **kwargs)

Generate a response using OpenAI's API.

Parameters:

Name Type Description Default
prompt str

The input prompt for the model.

 required
**kwargs

Additional parameters for the OpenAI API call, such as: - model (str): Model name to use (default: self.model_name). - max_tokens (int): Maximum number of tokens to generate (default: 128). - temperature (float): Sampling temperature (default: 0.7).

 {}

Returns:

Name Type Description
str str

The generated response as a string.
Notes
  • Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
  • All generation parameters can be overridden via kwargs.
Example 
answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
Source code in rankify/generator/models/openai_model.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def generate(self, prompt: str, **kwargs) -> str:
    """
    Generate a response using OpenAI's API.

    Args:
        prompt (str): The input prompt for the model.
        **kwargs: Additional parameters for the OpenAI API call, such as:
            - model (str): Model name to use (default: self.model_name).
            - max_tokens (int): Maximum number of tokens to generate (default: 128).
            - temperature (float): Sampling temperature (default: 0.7).

    Returns:
        str: The generated response as a string.

    Notes:
        - Default parameters: model=self.model_name, max_tokens=128, temperature=0.7.
        - All generation parameters can be overridden via `kwargs`.

    Example:
        ```python
        answer = model.generate("What is the capital of France?", max_tokens=64, temperature=0.5)
        ```
    """
    # Set default parameters for the OpenAI API call
    kwargs.setdefault("model", self.model_name)
    kwargs.setdefault("max_tokens", 128)
    kwargs.setdefault("temperature", 0.7)

    # Call the OpenAI API using the OpenaiClient
    response = self.client.chat(messages=[{"role": "user", "content": prompt}], return_text=True, **kwargs)
    return response

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 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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 
23
24
25
26
@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 
28
29
30
def embed(self, text: str, **kwargs) -> List[float]:
    """Optional: Generate embeddings for the given text."""
    raise NotImplementedError("Embedding is not required for this implementation.")

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 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
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 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
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))

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 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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 
23
24
25
26
27
28
29
30
31
32
33
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 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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

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 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
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 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
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)

load_model(model_name, **kwargs)

Source code in rankify/utils/generator/huggingface_models/model_utils.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def load_model(model_name, **kwargs):
    model_parallelism = kwargs.get("model_parallelism", False)
    cache_dir = kwargs.get("cache_dir", None)
    auth_token = kwargs.get("auth_token", None)
    torch_dtype = kwargs.get("torch_dtype", None)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    device_count = torch.cuda.device_count()

    config = AutoConfig.from_pretrained(model_name)
    model_args = {}
    if cache_dir is not None:
        model_args["cache_dir"] = cache_dir
    if model_parallelism:
        model_args["device_map"] = "auto"
        model_args["low_cpu_mem_usage"] = True
    if hasattr(config, "torch_dtype") and config.torch_dtype is not None:
        model_args["torch_dtype"] = config.torch_dtype
    if torch_dtype is not None:
        model_args["torch_dtype"] = torch_dtype  # overload dtype if user specifies
    if auth_token is not None:
        model_args["use_auth_token"] = auth_token

    model = AutoModelForCausalLM.from_pretrained(model_name, **model_args).eval()
    if not model_parallelism:
        model = model.to(device)

    return model

load_tokenizer(model_name)

Source code in rankify/utils/generator/huggingface_models/model_utils.py 
6
7
def load_tokenizer(model_name):
    return AutoTokenizer.from_pretrained(model_name)

model_factory(model_name, backend, method, stop_at_period=False, **kwargs)

Factory function to instantiate and return the appropriate RAG model based on the backend.

Parameters:

Name Type Description Default
model_name str

Name of the model to use (e.g., 'meta-llama/Meta-Llama-3.1-8B', 'nq_reader_base').

 required
backend str

Backend type ('openai', 'huggingface', 'fid', 'litellm', 'vllm').

 required
method str

RAG method or prompt strategy (e.g., 'zero-shot', 'fid').

 required
stop_at_period bool

If True, truncate output at first period (default: False).

 False
**kwargs

Additional model-specific parameters, such as API keys or generation settings.

 {}

Returns:

Name Type Description
BaseRAGModel BaseRAGModel

An instance of the selected model class, ready for text generation.

Raises:

Type Description
ValueError

If an unsupported backend is specified.
Notes
  • Automatically initializes the required prompt generator.
  • For 'openai', expects 'api_keys' in kwargs.
  • For 'huggingface', loads tokenizer and model locally.
  • For 'fid', instantiates Fusion-in-Decoder model.
  • For 'litellm', expects 'api_key' in kwargs.
  • For 'vllm', passes all kwargs to VLLMModel.
Example 
model = model_factory(
    model_name="meta-llama/Meta-Llama-3.1-8B",
    backend="huggingface",
    method="zero-shot",
    torch_dtype=torch.float16
)
Source code in rankify/generator/models/model_factory.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def model_factory(model_name: str, backend: str, method: str, stop_at_period=False, **kwargs) -> BaseRAGModel:
    """
    Factory function to instantiate and return the appropriate RAG model based on the backend.

    Args:
        model_name (str): Name of the model to use (e.g., 'meta-llama/Meta-Llama-3.1-8B', 'nq_reader_base').
        backend (str): Backend type ('openai', 'huggingface', 'fid', 'litellm', 'vllm').
        method (str): RAG method or prompt strategy (e.g., 'zero-shot', 'fid').
        stop_at_period (bool, optional): If True, truncate output at first period (default: False).
        **kwargs: Additional model-specific parameters, such as API keys or generation settings.

    Returns:
        BaseRAGModel: An instance of the selected model class, ready for text generation.

    Raises:
        ValueError: If an unsupported backend is specified.

    Notes:
        - Automatically initializes the required prompt generator.
        - For 'openai', expects 'api_keys' in kwargs.
        - For 'huggingface', loads tokenizer and model locally.
        - For 'fid', instantiates Fusion-in-Decoder model.
        - For 'litellm', expects 'api_key' in kwargs.
        - For 'vllm', passes all kwargs to VLLMModel.

    Example:
        ```python
        model = model_factory(
            model_name="meta-llama/Meta-Llama-3.1-8B",
            backend="huggingface",
            method="zero-shot",
            torch_dtype=torch.float16
        )
        ```
    """
    prompt_generator = PromptGenerator(model_type=model_name, method=method)
    if backend == "openai":
        return OpenAIModel(model_name, kwargs["api_keys"], prompt_generator,
                           base_url=kwargs.get("base_url", None))
    elif backend == "huggingface":
        tokenizer = load_tokenizer(model_name)
        model = load_model(model_name, **kwargs) 
        return HuggingFaceModel(model_name, tokenizer, model, prompt_generator, stop_at_period=stop_at_period)
    elif backend == "fid":
        return FiDModel(method, model_name, **kwargs)
    elif backend == "litellm":
        return LitellmModel(model_name, kwargs["api_key"], prompt_generator)
    elif backend == "vllm":
        return VLLMModel(model_name, prompt_generator, **kwargs)
    else:
        raise ValueError(f"Unsupported backend: {backend}")