logitorch.models.fld

Classes

FLDAllAtOnceProver

Initializes the FLDAllAtOnceProver model.

Module Contents

class logitorch.models.fld.FLDAllAtOnceProver(pretrained_t5_model: str)

Bases: torch.nn.Module

Initializes the FLDAllAtOnceProver model.

Args:

pretrained_t5_model (str): The name or path of the pretrained T5 model.

forward(x: Dict[str, torch.Tensor], y: torch.Tensor = None) → transformers.modeling_outputs.SequenceClassifierOutput

Performs a forward pass of the model.

Args:

x (Dict[str, torch.Tensor]): The input tensors. y (torch.Tensor, optional): The labels tensor. Defaults to None.

Returns:

SequenceClassifierOutput: The output of the model.

predict(prompt: str, num_beams: int = 5, max_length: int = 1000, device: str = 'cpu') → List[str]

Generates predictions based on the given prompt.

Args:

prompt (str): The input prompt. num_beams (int, optional): The number of beams for beam search. Defaults to 5. max_length (int, optional): The maximum length of the generated sequence. Defaults to 1000. device (str, optional): The device to run the model on. Defaults to “cpu”.

Returns:

List[str]: The generated predictions.

model

tokenizer