logitorch.models.proofwriter
============================

.. py:module:: logitorch.models.proofwriter


Classes
-------

.. autoapisummary::

   logitorch.models.proofwriter.ProofWriter


Module Contents
---------------

.. py:class:: ProofWriter(pretrained_t5_model: str)

   Bases: :py:obj:`torch.nn.Module`


   A PyTorch module for generating proofs using the T5 model.

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

   Attributes:
       model (T5ForConditionalGeneration): The T5 model for proof generation.
       tokenizer (T5Tokenizer): The tokenizer for tokenizing input text.

   Methods:
       forward(x, y=None): Performs forward pass of the model.
       predict(context, question, num_beams=5, max_length=120, device="cpu"): Generates proof given context and question.

   Initializes the ProofWriter module.

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


   .. py:method:: forward(x: Dict[str, torch.Tensor], y: torch.Tensor = None) -> transformers.modeling_outputs.SequenceClassifierOutput

      Performs forward pass of the model.

      Args:
          | x (Dict[str, torch.Tensor]): The input tensors for the model.
          | y (torch.Tensor, optional): The labels for the model. Defaults to None.

      Returns:
          | SequenceClassifierOutput: The output of the model.



   .. py:method:: predict(context: str, question: str, num_beams: int = 5, max_length: int = 120, device: str = 'cpu') -> List[str]

      Generates proof given context and question.

      Args:
          | context (str): The context for proof generation.
          | question (str): The question for proof generation.
          | num_beams (int, optional): The number of beams for beam search. Defaults to 5.
          | max_length (int, optional): The maximum length of the generated proof. Defaults to 120.
          | device (str, optional): The device to run the model on. Defaults to "cpu".

      Returns:
          | List[str]: The generated proof.



   .. py:attribute:: model


   .. py:attribute:: tokenizer