gradio.py

"""
Main App
"""

import gradio as gr
from transformers import AutoModelForSeq2SeqLM

from src.tokenizers import IndoNLGTokenizer


tokenizer = IndoNLGTokenizer.from_pretrained("indobenchmark/indobart-v2")
model = AutoModelForSeq2SeqLM.from_pretrained("haryoaw/id-recigen-bart")


def predict_recipe(food: str) -> str:
    """
    Predict Ingredients Here!

    Parameters
    ----------
    food: str
        The food that will be used
    """
    inp = tokenizer(food, return_tensors="pt")["input_ids"]
    generated = model.generate(
        inp, max_length=500, do_sample=False, num_beams=10, num_beam_groups=2
    )
    returned_input: str = tokenizer.decode(generated[0], skip_special_tokens=True)
    returned_input = "\n".join([x.strip() for x in returned_input.split("||")])
    return returned_input


iface = gr.Interface(
    fn=predict_recipe,
    inputs=[gr.inputs.Textbox(placeholder="Food Name")],
    outputs="textbox",
)

if __name__ == "__main__":
    app, local_url, share_url = iface.launch(share=False)