import gradio as gr
from transformers import pipeline

title = "Question Answer on Text"
description = """
- Provide a piece of text.
- Ask a question or specify an entity you would like to extract.
- Receive an answer.

Try out the examples at the bottom.
"""


classifier = pipeline(
    "question-answering", model="deepset/roberta-base-squad2", model_max_length=512
)


def zero_shot_classification(text_input, question):
    prediction = classifier(
        context=text_input,
        question=question,
        truncation=True,
        max_length="max_length",
        padding=True,
    )
    return f'{prediction["answer"]}'


examples = [
    [
        """The invoice contains all the goods that are described in the order with number x9820be. 
    The goods are packaged and should arrive at the destination on 31/12/2023. 
    We kindly request to transfer the amount within 30 days of the invoice date to the account number IE12BOFI90000112345678. """,
        "what is the timeline to transfer the money?",
    ],
    [
        """The invoice contains all the goods that are described in the order with number x9820be. 
    The goods are packaged and should arrive at the destination on 31/12/2023. 
    We kindly request to transfer the amount within 30 days of the invoice date to the account number IE12BOFI90000112345678""",
        "what is the account number to transfer the money?",
    ],[
        """The invoice contains all the goods that are described in the order with number x9820be. 
    The goods are packaged and should arrive at the destination on 31/12/2023. 
    We kindly request to transfer the amount within 30 days of the invoice date to the account number IE12BOFI90000112345678""",
        "order number",
    ],
    [
        """By signing below, Shipper hereby declares that the contents of this consignment are fully and accurately described above by the proper shipping name and are classified,
packaged, marked and labelled/placarded, and are in all respects in proper condition for transport according to applicable governmental regulations. As shipper, I hereby
certify that the liquid industrial by-product(s) are fully and accurately described on this shipping document, in proper condition for transport, and that the information
contained on the shipping document is factual.""",
        "are the goods described?",
    ],
]

gr.Interface(
    fn=zero_shot_classification,
    inputs=[
        gr.inputs.Textbox(lines=10, label="Text", placeholder="Paste text here ..."),
        gr.inputs.Textbox(
            lines=2,
            label="Question",
            placeholder="Ask what you want to retrieve from the vacancy.",
        ),
    ],
    outputs=gr.outputs.Textbox(label="Answer"),
    title=title,
    description=description,
    examples=examples,
).launch(share=False)