Skip to content

Generator

rankify.generator.generator

RAG_METHODS = {'in-context-ralm': InContextRALMRAG, 'fid': FiDRAGMethod, 'zero-shot': ZeroShotRAG, 'basic-rag': BasicRAG, 'chain-of-thought-rag': ChainOfThoughtRAG, 'self-consistency-rag': SelfConsistencyRAG, 'react-rag': ReActRAG} module-attribute

BasicRAG

Bases: BaseRAGMethod

BasicRAG (Naive RAG) Method for Retrieval-Augmented Generation.

Implements a simple RAG technique that answers questions by concatenating retrieved contexts and passing them, along with the question, to the underlying model. This is the most straightforward RAG approach.

Attributes:

Name Type Description
model BaseRAGModel

The RAG model instance used for generation.
References
  • Lewis et al. Retrieval-augmented generation for knowledge-intensive nlp tasks**
    Paper
Example 
from rankify.dataset.dataset import Document, Question, Answer, Context
from rankify.generator.generator import Generator

# Sample question and contexts
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)
]

# Create a Document
doc = Document(question=question, answers= answers, contexts=contexts)

# Initialize Generator (e.g., Meta Llama, with huggingface backend)
generator = Generator(method="basic-rag", model_name='meta-llama/Meta-Llama-3.1-8B-Instruct', backend="huggingface")

# Generate answer
generated_answers = generator.generate([doc])
print(generated_answers)  # Output: ["Paris"]
Notes
  • This method does not apply advanced reasoning or fusion techniques.
  • Suitable as a baseline for comparison with more sophisticated RAG methods.
Source code in rankify/generator/rag_methods/basic_rag.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
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
class BasicRAG(BaseRAGMethod):
    """
    **BasicRAG (Naive RAG) Method** for Retrieval-Augmented Generation.

    Implements a simple RAG technique that answers questions by concatenating retrieved contexts and passing them,
    along with the question, to the underlying model. This is the most straightforward RAG approach.

    Attributes:
        model (BaseRAGModel): The RAG model instance used for generation.

    References:
        - **Lewis et al. **Retrieval-augmented generation for knowledge-intensive nlp tasks**  
          [Paper](https://proceedings.neurips.cc/paper/2020/hash/6b493230205f780e1bc26945df7481e5-Abstract.html)

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

        # Sample question and contexts
        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)
        ]

        # Create a Document
        doc = Document(question=question, answers= answers, contexts=contexts)

        # Initialize Generator (e.g., Meta Llama, with huggingface backend)
        generator = Generator(method="basic-rag", model_name='meta-llama/Meta-Llama-3.1-8B-Instruct', backend="huggingface")

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

    Notes:
        - This method does not apply advanced reasoning or fusion techniques.
        - Suitable as a baseline for comparison with more sophisticated RAG methods.
    """
    def __init__(self, model: BaseRAGModel, **kwargs):
        """
        Initialize the BasicRAG method.

        Args:
            model (BaseRAGModel): The RAG model instance used for generation.
        """
        super().__init__(model=model)

    def answer_questions(self, documents: List[Document], custom_prompt=None, **kwargs) -> List[str]:
        """
        Answer questions for a list of documents using the model.

        Args:
            documents (List[Document]): A list of Document objects containing questions and contexts.
            custom_prompt (str, optional): Custom prompt to override default prompt generation.
            **kwargs: Additional parameters for the model's generate method.

        Returns:
            List[str]: Answers generated for each document.

        Notes:
            - Concatenates all context passages and passes them with the question to the model.
            - Uses the model's prompt generator to construct prompts.
        """
        answers = []

        for document in tqdm(documents, desc="Answering questions", unit="q"):
            # Extract question and contexts from the document
            question = document.question.question
            contexts = [context.text for context in document.contexts]

            # Construct the prompt
            prompt = self.model.prompt_generator.generate_user_prompt(question, contexts, custom_prompt)
            # Generate the answer using the model
            answer = self.model.generate(prompt=prompt, **kwargs)

            # Append the answer to the list
            answers.append(answer)
        return answers

__init__(model, **kwargs)

Initialize the BasicRAG method.

Parameters:

Name Type Description Default
model BaseRAGModel

The RAG model instance used for generation.

 required
Source code in rankify/generator/rag_methods/basic_rag.py 
51
52
53
54
55
56
57
58
def __init__(self, model: BaseRAGModel, **kwargs):
    """
    Initialize the BasicRAG method.

    Args:
        model (BaseRAGModel): The RAG model instance used for generation.
    """
    super().__init__(model=model)

answer_questions(documents, custom_prompt=None, **kwargs)

Answer questions for a list of documents using the model.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document objects containing questions and contexts.

 required
custom_prompt str

Custom prompt to override default prompt generation.

 None
**kwargs

Additional parameters for the model's generate method.

 {}

Returns:

Type Description
List[str]

List[str]: Answers generated for each document.
Notes
  • Concatenates all context passages and passes them with the question to the model.
  • Uses the model's prompt generator to construct prompts.
Source code in rankify/generator/rag_methods/basic_rag.py 
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
def answer_questions(self, documents: List[Document], custom_prompt=None, **kwargs) -> List[str]:
    """
    Answer questions for a list of documents using the model.

    Args:
        documents (List[Document]): A list of Document objects containing questions and contexts.
        custom_prompt (str, optional): Custom prompt to override default prompt generation.
        **kwargs: Additional parameters for the model's generate method.

    Returns:
        List[str]: Answers generated for each document.

    Notes:
        - Concatenates all context passages and passes them with the question to the model.
        - Uses the model's prompt generator to construct prompts.
    """
    answers = []

    for document in tqdm(documents, desc="Answering questions", unit="q"):
        # Extract question and contexts from the document
        question = document.question.question
        contexts = [context.text for context in document.contexts]

        # Construct the prompt
        prompt = self.model.prompt_generator.generate_user_prompt(question, contexts, custom_prompt)
        # Generate the answer using the model
        answer = self.model.generate(prompt=prompt, **kwargs)

        # Append the answer to the list
        answers.append(answer)
    return answers

ChainOfThoughtRAG

Bases: BaseRAGMethod

Chain-of-Thought RAG for Open-Domain Question Answering.

This class implements a chain-of-thought retrieval-augmented generation (RAG) method, where the model generates answers by reasoning step-by-step using the provided contexts.

Attributes:

Name Type Description
model BaseRAGModel

The underlying model used for text generation.
References
  • Wei et al. Chain-of-thought prompting elicits reasoning in large language models**
    Paper

Methods:

Name Description
answer_questions

List[Document], **kwargs) -> List[str]: Generates answers for a list of documents using chain-of-thought reasoning.
Example 
from rankify.dataset.dataset import Document, Question, Answer, Context
from rankify.generator.generator import Generator

# Sample question and contexts
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)
]

# Create a Document
doc = Document(question=question, answers= answers, contexts=contexts)

# Initialize Generator (e.g., Meta Llama, with huggingface backend)
generator = Generator(method="chain-of-thought-rag", model_name='meta-llama/Meta-Llama-3.1-8B-Instruct', backend="huggingface")

# Generate answer
generated_answers = generator.generate([doc])
print(generated_answers)  # Output: ["Paris"]
Notes
  • This method uses chain-of-thought reasoning to generate more detailed and logical answers.
  • It is particularly useful for complex questions that require multi-step reasoning.
Source code in rankify/generator/rag_methods/chain_of_thought_rag.py 
 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 ChainOfThoughtRAG(BaseRAGMethod):
    """
    **Chain-of-Thought RAG** for Open-Domain Question Answering.

    This class implements a chain-of-thought retrieval-augmented generation (RAG) method, 
    where the model generates answers by reasoning step-by-step using the provided contexts.

    Attributes:
        model (BaseRAGModel): The underlying model used for text generation.

    References:
        - **Wei et al. **Chain-of-thought prompting elicits reasoning in large language models**  
          [Paper](https://proceedings.neurips.cc/paper/2022/hash/9d5609613524ecf4f15af0f7b31abca4-Abstract-Conference.html)

    Methods:
        answer_questions(documents: List[Document], **kwargs) -> List[str]:
            Generates answers for a list of documents using chain-of-thought reasoning.

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

        # Sample question and contexts
        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)
        ]

        # Create a Document
        doc = Document(question=question, answers= answers, contexts=contexts)

        # Initialize Generator (e.g., Meta Llama, with huggingface backend)
        generator = Generator(method="chain-of-thought-rag", model_name='meta-llama/Meta-Llama-3.1-8B-Instruct', backend="huggingface")

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

    Notes:
        - This method uses chain-of-thought reasoning to generate more detailed and logical answers.
        - It is particularly useful for complex questions that require multi-step reasoning.

    """
    def __init__(self, model: BaseRAGModel, **kwargs):
        super().__init__(model=model)

    def answer_questions(self, documents: List[Document], custom_prompt=None, **kwargs) -> List[str]:
        """
        Answer questions for a list of documents using chain-of-thought reasoning.

        Args:
            documents (List[Document]): A list of Document objects containing questions and contexts.
            custom_prompt (str, optional): Custom prompt to override default prompt generation.
            **kwargs: Additional parameters for the model's generate method.

        Returns:
            List[str]: Answers generated for each document, with chain-of-thought reasoning.

        Notes:
            - Uses the model's prompt generator to build chain-of-thought prompts from question and contexts.
            - Suitable for complex questions requiring multi-step inference.
        """
        answers = []

        for document in tqdm(documents, desc="Answering questions", unit="q"):
            # Extract question and contexts from the document
            question = document.question.question
            contexts = [context.text for context in document.contexts]

            # Construct the prompt
            prompt = self.model.prompt_generator.generate_user_prompt(question, contexts, custom_prompt)

            # Generate the answer using the model
            answer = self.model.generate(prompt=prompt, **kwargs)

            # Append the answer to the list
            answers.append(answer)
        return answers

answer_questions(documents, custom_prompt=None, **kwargs)

Answer questions for a list of documents using chain-of-thought reasoning.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document objects containing questions and contexts.

 required
custom_prompt str

Custom prompt to override default prompt generation.

 None
**kwargs

Additional parameters for the model's generate method.

 {}

Returns:

Type Description
List[str]

List[str]: Answers generated for each document, with chain-of-thought reasoning.
Notes
  • Uses the model's prompt generator to build chain-of-thought prompts from question and contexts.
  • Suitable for complex questions requiring multi-step inference.
Source code in rankify/generator/rag_methods/chain_of_thought_rag.py 
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
def answer_questions(self, documents: List[Document], custom_prompt=None, **kwargs) -> List[str]:
    """
    Answer questions for a list of documents using chain-of-thought reasoning.

    Args:
        documents (List[Document]): A list of Document objects containing questions and contexts.
        custom_prompt (str, optional): Custom prompt to override default prompt generation.
        **kwargs: Additional parameters for the model's generate method.

    Returns:
        List[str]: Answers generated for each document, with chain-of-thought reasoning.

    Notes:
        - Uses the model's prompt generator to build chain-of-thought prompts from question and contexts.
        - Suitable for complex questions requiring multi-step inference.
    """
    answers = []

    for document in tqdm(documents, desc="Answering questions", unit="q"):
        # Extract question and contexts from the document
        question = document.question.question
        contexts = [context.text for context in document.contexts]

        # Construct the prompt
        prompt = self.model.prompt_generator.generate_user_prompt(question, contexts, custom_prompt)

        # Generate the answer using the model
        answer = self.model.generate(prompt=prompt, **kwargs)

        # Append the answer to the list
        answers.append(answer)
    return answers

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))

FiDRAGMethod

Bases: BaseRAGMethod

FiD RAG Method for Open-Domain Question Answering.

This class implements a retrieval-augmented generation (RAG) method using the Fusion-in-Decoder (FiD) approach. The FiD model aggregates information from multiple retrieved passages to generate context-aware answers.

References
  • Izacard & Grave Leveraging Passage Retrieval with Generative Models for Open-Domain QA
    Paper

Attributes:

Name Type Description
model BaseRAGModel

The underlying FiD model used for text generation.

Methods:

Name Description
answer_questions

List[Document], **kwargs) -> List[str]: Generates answers for a list of documents using the FiD model.
See Also
  • FiDModel: Class for FiDModel, containing the FiD specific logic.
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
  • This class was created to keep the unified interface of RAG methods.
  • Since FiD is a specific RAG technique that relies on the full transformer architecture, the logic is included in the model, see
Source code in rankify/generator/rag_methods/fid_rag_method.py 
 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
class FiDRAGMethod(BaseRAGMethod):
    """
    **FiD RAG Method** for Open-Domain Question Answering.

    This class implements a retrieval-augmented generation (RAG) method using the 
    Fusion-in-Decoder (FiD) approach. The FiD model aggregates information from 
    multiple retrieved passages to generate context-aware answers.

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

    Attributes:
        model (BaseRAGModel): The underlying FiD model used for text generation.

    Methods:
        answer_questions(documents: List[Document], **kwargs) -> List[str]:
            Generates answers for a list of documents using the FiD model.

    See Also:
        - `FiDModel`: Class for FiDModel, containing the FiD specific logic.

    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:
        - This class was created to keep the unified interface of RAG methods.
        - Since FiD is a specific RAG technique that relies on the full transformer architecture,
          the logic is included in the model, see 

    """
    def __init__(self, model: BaseRAGModel, **kwargs):
        super().__init__(model=model)

    def answer_questions(self, documents: List[Document], custom_prompt=None, **kwargs) -> List[str]:
        """
        Answer questions for a list of documents using the FiDModel.

        Args:
            documents (List[Document]): A list of Document objects containing questions and contexts.

        Returns:
            Lists[str]: An answer based on the given documents and question.
        """

        # Generate the answer using the model
        answer = self.model.generate(documents, **kwargs)

        return answer

answer_questions(documents, custom_prompt=None, **kwargs)

Answer questions for a list of documents using the FiDModel.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document objects containing questions and contexts.

 required

Returns:

Type Description
List[str]

Lists[str]: An answer based on the given documents and question.
Source code in rankify/generator/rag_methods/fid_rag_method.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def answer_questions(self, documents: List[Document], custom_prompt=None, **kwargs) -> List[str]:
    """
    Answer questions for a list of documents using the FiDModel.

    Args:
        documents (List[Document]): A list of Document objects containing questions and contexts.

    Returns:
        Lists[str]: An answer based on the given documents and question.
    """

    # Generate the answer using the model
    answer = self.model.generate(documents, **kwargs)

    return answer

InContextRALMRAG

Bases: BaseRAGMethod

In-Context Retrieval-Augmented Language Model (RALM) Generator.

This class implements In-Context Retrieval-Augmented Language Models (RALM) for context-aware text generation. It integrates retrieved passages into the input prompt for few-shot learning, improving response quality.

Attributes:

Name Type Description
model BaseRAGModel

The underlying RAG model used for generation.
tokenizer

Tokenizer associated with the model.
device str

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

Directory to store downloaded models and cache files.
num_docs int

Number of retrieved contexts to include in the prompt (default: 1).
max_length int

Maximum number of tokens the model can process.
max_tokens_to_generate int

Maximum number of tokens the model generates in response.

Methods:

Name Description
_build_qa_prompt

Constructs a QA prompt with retrieved passages for in-context learning.
_prepare_dataloader

Converts Document objects into a dataset formatted for RALM.
answer_questions

Generates answers for a list of documents using RALM.
References
  • Ram et al. In-Context Retrieval-Augmented Language Models
    Paper
See Also
  • BaseRAGMethod: Parent class for RAG techniques.
  • Few-Shot Learning: This method leverages retrieved passages for in-context QA.
Example 
from rankify.dataset.dataset import Document, Question, Answer, Context
from rankify.generator.generator import Generator

# Sample question and contexts
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)
]

# Create a Document
doc = Document(question=question, answers= answers, contexts=contexts)

# Initialize Generator (e.g., Meta Llama, with huggingface backend)
generator = Generator(method="in-context-ralm", model_name='meta-llama/Meta-Llama-3.1-8B-Instruct', backend="huggingface")

# Generate answer
generated_answers = generator.generate([doc])
print(generated_answers)  # Output: ["Paris"]
Notes
  • RALM dynamically integrates retrieved passages to guide response generation.
  • Optimized for retrieval-augmented question answering (RAG).
  • Prompts are constructed to encourage factual, concise answers suitable for evaluation and comparison.
Source code in rankify/generator/rag_methods/in_context_ralm_rag.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
 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
class InContextRALMRAG(BaseRAGMethod):
    """
    **In-Context Retrieval-Augmented Language Model (RALM) Generator**.


    This class implements **In-Context Retrieval-Augmented Language Models (RALM)** for **context-aware text generation**.
    It **integrates retrieved passages** into the input prompt for **few-shot learning**, improving response quality.

    Attributes:
        model (BaseRAGModel): The underlying RAG model used for generation.
        tokenizer: Tokenizer associated with the model.
        device (str): Device used for inference (`'cuda'` or `'cpu'`).
        cache_dir (str): Directory to store downloaded models and cache files.
        num_docs (int): Number of retrieved contexts to include in the prompt (default: 1).
        max_length (int): Maximum number of tokens the model can process.
        max_tokens_to_generate (int): Maximum number of tokens the model generates in response.

    Methods:
        _build_qa_prompt(example): Constructs a QA prompt with retrieved passages for in-context learning.
        _prepare_dataloader(documents): Converts Document objects into a dataset formatted for RALM.
        answer_questions(documents, custom_prompt=None): Generates answers for a list of documents using RALM.

    References:
        - **Ram et al.** *In-Context Retrieval-Augmented Language Models*  
          [Paper](https://arxiv.org/abs/2302.00083)

    See Also:
        - `BaseRAGMethod`: Parent class for RAG techniques.
        - Few-Shot Learning: This method leverages retrieved passages for in-context QA.

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

        # Sample question and contexts
        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)
        ]

        # Create a Document
        doc = Document(question=question, answers= answers, contexts=contexts)

        # Initialize Generator (e.g., Meta Llama, with huggingface backend)
        generator = Generator(method="in-context-ralm", model_name='meta-llama/Meta-Llama-3.1-8B-Instruct', backend="huggingface")

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

    Notes:
        - RALM dynamically integrates retrieved passages to guide response generation.
        - Optimized for retrieval-augmented question answering (RAG).
        - Prompts are constructed to encourage factual, concise answers suitable for evaluation and comparison.
    """

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

    def __init__(self, model:BaseRAGModel, **kwargs):
        """
        Initializes the RALM Generator.

        Args:
            method (str): The generator type (`"in-context-ralm"`).
            model_name (str): The name of the pre-trained RALM model (e.g., `"meta-llama/Llama-3.1-8B"`).
            **kwargs: Additional parameters for model configuration.

        Example:
            ```python
            generator = InContextRALMGenerator(method="in-context-ralm", model_name="meta-llama/Llama-3.1-8B")
            ```
        """
        self.base_rag_model = model
        self.model = model.model
        self.tokenizer = model.tokenizer
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.cache_dir = self.CACHE_DIR
        self.num_docs = kwargs.get("num_docs", 1)  # Default: 1 supporting document



        self.tokenizer.pad_token = self.tokenizer.eos_token

        self.config = AutoConfig.from_pretrained(model.model_name)

        # Set generation parameters
        self.max_length = self.config.n_positions if hasattr(self.config, "n_positions") else self.config.max_position_embeddings
        self.max_tokens_to_generate = kwargs.get("max_tokens_to_generate", 10)


    def _prepare_dataloader(self, documents: list[Document]):
        """
        Converts `Document` objects into a dataset **formatted for RALM**.

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

        Returns:
            list: A list of dictionaries formatted for in-context learning.

        Example:
            ```python
            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)

        return examples

    def answer_questions(self, documents: List[Document], custom_prompt=None) -> List[str]:
        """
        Generates answers for **a list of documents** using RALM.

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

        Returns:
            List[str]: A list of generated answers.
        """
        eval_dataset = self._prepare_dataloader(documents)

        results = []
        for example in tqdm(eval_dataset, desc="Answering questions", unit="q"):
            context_strs = [f"{ctx['title']}\n\n{ctx['text']}" for ctx in example["ctxs"][:self.num_docs]]
            prompt = self.base_rag_model.prompt_generator.generate_user_prompt(
                question=example["question"],
                contexts=context_strs,
                custom_prompt=custom_prompt
            )

            tokenized_input = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
            input_ids = tokenized_input.input_ids.to(self.device)
            attention_mask = tokenized_input.attention_mask.to(self.device)  # Extract attention mask

            if input_ids.shape[-1] > self.max_length - self.max_tokens_to_generate:
                input_ids = input_ids[..., -(self.max_length - self.max_tokens_to_generate):]
                attention_mask = attention_mask[..., -(self.max_length - self.max_tokens_to_generate):]

            with torch.no_grad():
                outputs = self.model.generate(input_ids, attention_mask=attention_mask, max_new_tokens=self.max_tokens_to_generate,  pad_token_id=self.tokenizer.pad_token_id )

            # Extract generated text
            generation_str = self.tokenizer.decode(outputs[0].cpu(), skip_special_tokens=True)
            answer = generation_str[len(prompt):].split("\n")[0]

            results.append(answer)

        return results

__init__(model, **kwargs)

Initializes the RALM Generator.

Parameters:

Name Type Description Default
method str

The generator type ("in-context-ralm").

 required
model_name str

The name of the pre-trained RALM model (e.g., "meta-llama/Llama-3.1-8B").

 required
**kwargs

Additional parameters for model configuration.

 {}
Example 
generator = InContextRALMGenerator(method="in-context-ralm", model_name="meta-llama/Llama-3.1-8B")
Source code in rankify/generator/rag_methods/in_context_ralm_rag.py 
 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
def __init__(self, model:BaseRAGModel, **kwargs):
    """
    Initializes the RALM Generator.

    Args:
        method (str): The generator type (`"in-context-ralm"`).
        model_name (str): The name of the pre-trained RALM model (e.g., `"meta-llama/Llama-3.1-8B"`).
        **kwargs: Additional parameters for model configuration.

    Example:
        ```python
        generator = InContextRALMGenerator(method="in-context-ralm", model_name="meta-llama/Llama-3.1-8B")
        ```
    """
    self.base_rag_model = model
    self.model = model.model
    self.tokenizer = model.tokenizer
    self.device = "cuda" if torch.cuda.is_available() else "cpu"
    self.cache_dir = self.CACHE_DIR
    self.num_docs = kwargs.get("num_docs", 1)  # Default: 1 supporting document



    self.tokenizer.pad_token = self.tokenizer.eos_token

    self.config = AutoConfig.from_pretrained(model.model_name)

    # Set generation parameters
    self.max_length = self.config.n_positions if hasattr(self.config, "n_positions") else self.config.max_position_embeddings
    self.max_tokens_to_generate = kwargs.get("max_tokens_to_generate", 10)

answer_questions(documents, custom_prompt=None)

Generates answers for a list of documents using RALM.

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.
Source code in rankify/generator/rag_methods/in_context_ralm_rag.py 
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
def answer_questions(self, documents: List[Document], custom_prompt=None) -> List[str]:
    """
    Generates answers for **a list of documents** using RALM.

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

    Returns:
        List[str]: A list of generated answers.
    """
    eval_dataset = self._prepare_dataloader(documents)

    results = []
    for example in tqdm(eval_dataset, desc="Answering questions", unit="q"):
        context_strs = [f"{ctx['title']}\n\n{ctx['text']}" for ctx in example["ctxs"][:self.num_docs]]
        prompt = self.base_rag_model.prompt_generator.generate_user_prompt(
            question=example["question"],
            contexts=context_strs,
            custom_prompt=custom_prompt
        )

        tokenized_input = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
        input_ids = tokenized_input.input_ids.to(self.device)
        attention_mask = tokenized_input.attention_mask.to(self.device)  # Extract attention mask

        if input_ids.shape[-1] > self.max_length - self.max_tokens_to_generate:
            input_ids = input_ids[..., -(self.max_length - self.max_tokens_to_generate):]
            attention_mask = attention_mask[..., -(self.max_length - self.max_tokens_to_generate):]

        with torch.no_grad():
            outputs = self.model.generate(input_ids, attention_mask=attention_mask, max_new_tokens=self.max_tokens_to_generate,  pad_token_id=self.tokenizer.pad_token_id )

        # Extract generated text
        generation_str = self.tokenizer.decode(outputs[0].cpu(), skip_special_tokens=True)
        answer = generation_str[len(prompt):].split("\n")[0]

        results.append(answer)

    return results

Generator

Generator Interface for Retrieval-Augmented Generation (RAG) techniques.

This class serves as a unified wrapper for all RAG methods and underlying model endpoints in Rankify. It dynamically initializes the appropriate model and RAG method, using the model_factory method, based on the provided method, model_name, and backend.

Attributes:

Name Type Description
rag_method BaseRAGMethod

The initialized RAG method instance, combining the selected model and technique.

Raises:

Type Description
ValueError

If the specified method is not available in RAG_METHODS.
Example 
from rankify.generator.generator import Generator
from rankify.dataset.dataset import Document, Question

# Example: Fusion-in-Decoder (FiD)
generator = Generator(method="fid", model_name="nq_reader_base", backend="fid")
documents = [Document(question=Question("Who wrote Hamlet?"))]
generated_answers = generator.generate(documents)
print(generated_answers[0])  # Output: "William Shakespeare wrote Hamlet in the early 1600s."

# Example: Chain-of-Thought RAG
generator = Generator(method="chain-of-thought-rag", model_name="meta-llama/Meta-Llama-3.1-8B-Instruct", backend="huggingface")
documents = [Document(question=Question("What is the capital of France?"))]
print(generator.generate(documents)[0])  # Output: "Paris"
Notes
  • The generator puts together both the model endpoint (e.g., OpenAI, HuggingFace, FiD, vLLM) and the RAG method (e.g., basic, chain-of-thought, self-consistency).
  • All configuration (model, backend, method, and additional kwargs) is passed to the appropriate factory and method class.
  • Use generate() to obtain answers for a batch of documents, calling answer_questions() of the specified RAG method.
Source code in rankify/generator/generator.py 
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
class Generator:
    """
    **Generator Interface** for Retrieval-Augmented Generation (RAG) techniques.

    This class serves as a unified wrapper for all RAG methods and underlying model endpoints in Rankify.
    It dynamically initializes the appropriate model and RAG method, using the model_factory method,
    based on the provided `method`, `model_name`, and `backend`.

    Attributes:
        rag_method (BaseRAGMethod): The initialized RAG method instance, combining the selected model and technique.

    Raises:
        ValueError: If the specified `method` is not available in `RAG_METHODS`.

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

        # Example: Fusion-in-Decoder (FiD)
        generator = Generator(method="fid", model_name="nq_reader_base", backend="fid")
        documents = [Document(question=Question("Who wrote Hamlet?"))]
        generated_answers = generator.generate(documents)
        print(generated_answers[0])  # Output: "William Shakespeare wrote Hamlet in the early 1600s."

        # Example: Chain-of-Thought RAG
        generator = Generator(method="chain-of-thought-rag", model_name="meta-llama/Meta-Llama-3.1-8B-Instruct", backend="huggingface")
        documents = [Document(question=Question("What is the capital of France?"))]
        print(generator.generate(documents)[0])  # Output: "Paris"
        ```

    Notes:
        - The generator puts together both the model endpoint (e.g., OpenAI, HuggingFace, FiD, vLLM) and the RAG method (e.g., basic, chain-of-thought, self-consistency).
        - All configuration (model, backend, method, and additional kwargs) is passed to the appropriate factory and method class.
        - Use `generate()` to obtain answers for a batch of documents, calling `answer_questions()` of the specified RAG method.
    """

    def __init__(self, method: str, model_name: str, backend:str = "huggingface", **kwargs):
        """
        Initializes the selected RAG method and model.

        Args:
            method (str): The RAG technique (e.g., "fid", "chain-of-thought-rag", "self-consistency-rag").
            model_name (str): The specific model name (e.g., "nq_reader_base", "meta-llama/Meta-Llama-3.1-8B-Instruct").
            backend (str): The model backend ("huggingface", "openai", "fid", "vllm", etc.).
            **kwargs: Additional parameters for model and method initialization.

        Raises:
            ValueError: If the specified `method` is not available in `RAG_METHODS`.
        """
        if method not in RAG_METHODS:
            raise ValueError(f"Generator method {method} is not supported.")

        # Initialize the generator model based on the method
        model = model_factory(model_name=model_name, backend=backend, method=method, **kwargs)

        # get the class for the specified method
        rag_method_class = RAG_METHODS[method]

        # Initialize the generator with the model and any additional parameters
        self.rag_method = rag_method_class(model, **kwargs)


    def generate(self, documents, custom_prompt=None,**kwargs):
        """
        Generates answers for a batch of input documents using the selected RAG method.

        Args:
            documents (List[Document]): A list of `Document` objects containing queries and retrieved contexts.
            custom_prompt (str, optional): Custom prompt to override default prompt generation.
            **kwargs: Additional parameters for the RAG method's answer logic or model generation.

        Returns:
            List[str]: A list of generated answers, one for each document.

        """
        return self.rag_method.answer_questions(documents, custom_prompt, **kwargs)

__init__(method, model_name, backend='huggingface', **kwargs)

Initializes the selected RAG method and model.

Parameters:

Name Type Description Default
method str

The RAG technique (e.g., "fid", "chain-of-thought-rag", "self-consistency-rag").

 required
model_name str

The specific model name (e.g., "nq_reader_base", "meta-llama/Meta-Llama-3.1-8B-Instruct").

 required
backend str

The model backend ("huggingface", "openai", "fid", "vllm", etc.).

 'huggingface'
**kwargs

Additional parameters for model and method initialization.

 {}

Raises:

Type Description
ValueError

If the specified method is not available in RAG_METHODS.
Source code in rankify/generator/generator.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def __init__(self, method: str, model_name: str, backend:str = "huggingface", **kwargs):
    """
    Initializes the selected RAG method and model.

    Args:
        method (str): The RAG technique (e.g., "fid", "chain-of-thought-rag", "self-consistency-rag").
        model_name (str): The specific model name (e.g., "nq_reader_base", "meta-llama/Meta-Llama-3.1-8B-Instruct").
        backend (str): The model backend ("huggingface", "openai", "fid", "vllm", etc.).
        **kwargs: Additional parameters for model and method initialization.

    Raises:
        ValueError: If the specified `method` is not available in `RAG_METHODS`.
    """
    if method not in RAG_METHODS:
        raise ValueError(f"Generator method {method} is not supported.")

    # Initialize the generator model based on the method
    model = model_factory(model_name=model_name, backend=backend, method=method, **kwargs)

    # get the class for the specified method
    rag_method_class = RAG_METHODS[method]

    # Initialize the generator with the model and any additional parameters
    self.rag_method = rag_method_class(model, **kwargs)

generate(documents, custom_prompt=None, **kwargs)

Generates answers for a batch of input documents using the selected RAG method.

Parameters:

Name Type Description Default
documents List[Document]

A list of Document objects containing queries and retrieved contexts.

 required
custom_prompt str

Custom prompt to override default prompt generation.

 None
**kwargs

Additional parameters for the RAG method's answer logic or model generation.

 {}

Returns:

Type Description

List[str]: A list of generated answers, one for each document.
Source code in rankify/generator/generator.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def generate(self, documents, custom_prompt=None,**kwargs):
    """
    Generates answers for a batch of input documents using the selected RAG method.

    Args:
        documents (List[Document]): A list of `Document` objects containing queries and retrieved contexts.
        custom_prompt (str, optional): Custom prompt to override default prompt generation.
        **kwargs: Additional parameters for the RAG method's answer logic or model generation.

    Returns:
        List[str]: A list of generated answers, one for each document.

    """
    return self.rag_method.answer_questions(documents, custom_prompt, **kwargs)

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}")