Generators

Base Generator

BaseGenerator

Bases: BaseModel, ABC

Base Generator Class.

Source code in src/fed_rag/base/generator.py

class BaseGenerator(BaseModel, ABC):
    """Base Generator Class."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @abstractmethod
    def generate(
        self,
        query: str | list[str] | Query | list[Query],
        context: str | list[str] | Context | list[Context],
        **kwargs: dict,
    ) -> str | list[str]:
        """Generate an output from a given query and context."""

    @abstractmethod
    def complete(
        self, prompt: str | list[str] | Prompt | list[Prompt], **kwargs: dict
    ) -> str | list[str]:
        """Completion interface for generator LLMs."""

    @property
    @abstractmethod
    def model(self) -> torch.nn.Module:
        """Model associated with this generator."""

    @property
    @abstractmethod
    def tokenizer(self) -> BaseTokenizer:
        """Tokenizer associated with this generator."""

    @abstractmethod
    def compute_target_sequence_proba(
        self, prompt: str | Prompt, target: str
    ) -> torch.Tensor:
        """Compute P(target | prompt).

        NOTE: this is used in LM Supervised Retriever fine-tuning.
        """

    @property
    @abstractmethod
    def prompt_template(self) -> str:
        """Prompt template for formating query and context."""

    @prompt_template.setter
    @abstractmethod
    def prompt_template(self, value: str) -> None:
        """Prompt template setter."""

model abstractmethod property

model

Model associated with this generator.

tokenizer abstractmethod property

tokenizer

Tokenizer associated with this generator.

prompt_template abstractmethod property writable

prompt_template

Prompt template for formating query and context.

generate abstractmethod

generate(query, context, **kwargs)

Generate an output from a given query and context.

Source code in src/fed_rag/base/generator.py

@abstractmethod
def generate(
    self,
    query: str | list[str] | Query | list[Query],
    context: str | list[str] | Context | list[Context],
    **kwargs: dict,
) -> str | list[str]:
    """Generate an output from a given query and context."""

complete abstractmethod

complete(prompt, **kwargs)

Completion interface for generator LLMs.

Source code in src/fed_rag/base/generator.py

@abstractmethod
def complete(
    self, prompt: str | list[str] | Prompt | list[Prompt], **kwargs: dict
) -> str | list[str]:
    """Completion interface for generator LLMs."""

compute_target_sequence_proba abstractmethod

compute_target_sequence_proba(prompt, target)

Compute P(target | prompt).

NOTE: this is used in LM Supervised Retriever fine-tuning.

Source code in src/fed_rag/base/generator.py

@abstractmethod
def compute_target_sequence_proba(
    self, prompt: str | Prompt, target: str
) -> torch.Tensor:
    """Compute P(target | prompt).

    NOTE: this is used in LM Supervised Retriever fine-tuning.
    """