app.py

"""
Main App
"""

import streamlit as st
from transformers import AutoModelForSeq2SeqLM

from src.tokenizers import IndoNLGTokenizer


@st.cache(allow_output_mutation=True)
def fetch_tokenizer_model():
    """
    Fetch tokenizer and model
    """
    tokenizer = IndoNLGTokenizer.from_pretrained("indobenchmark/indobart-v2")
    model = AutoModelForSeq2SeqLM.from_pretrained("haryoaw/id-recigen-bart")
    return tokenizer, model


tokenizer, model = fetch_tokenizer_model()


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

    Parameters
    ----------
    food: str
        The food that will be used

    Returns
    -------
    str
        Return the model here
    """
    inp = tokenizer(food.lower(), 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


def create_frontend() -> None:
    """
    Create front end streamlit here
    """
    st.markdown("# Food Ingredients Generator Indonesia Showcase!")
    st.write("🥑Generate your ingredients here!")
    
    with st.form("my_form"):
        food_name = st.text_input(
            "Food", value="Nasi Goreng Ayam", help="Input your food here!"
        )
        submitted = st.form_submit_button("Submit")
        if submitted:
            predicted = predict_recipe(food_name)
            st.markdown(f"## Bahan ( Ingredients ) `{food_name}`:")
            st.text(predicted)
    st.markdown("## Additional Note")
    st.write(
        "❗Please note that the model is trained with the food that use:"
    )
    for i, ingr in enumerate(("ayam", "tempe", "ikan", "kambing", "telur", "tahu", "sapi")):
        st.write(f"{i+1}. {ingr}")
    
    st.markdown("## Models")
    st.markdown(
        "🤗 Huggingface Model: [Link](https://huggingface.co/haryoaw/id-recigen-bart)"
    )
    st.write("Thank you 😊")


if __name__ == "__main__":
    create_frontend()