import streamlit as st
from transformers import pipeline
import time


@st.cache_resource
def load_pipeline(model_name):
    with st.spinner(f"Loading model {model_name}..."):
        return pipeline("text-generation", model=model_name)


MODEL = "teamapocalypseml/regben2ipa-byt5small"

MAPPER = {
    "byt5 small regional transcriber": "teamapocalypseml/regben2ipa-byt5small",
    "umt5 base regional transcriber": "teamapocalypseml/regben2ipa-umt5base",
    "mt5 base regional transcriber": "teamapocalypseml/regben2ipa-mt5-base",
}

st.title("Model configurations")
MODEL = st.selectbox("Select Model", [
    "byt5 small regional transcriber",
    "umt5 base regional transcriber",
    "mt5 base regional transcriber"
])
st.link_button(
    f"{MAPPER[MODEL]}", f"https://huggingface.co/{MAPPER[MODEL]}",
    type="tertiary",
    icon="🔗"
)

model = load_pipeline(MAPPER[MODEL])

model = pipeline("text2text-generation", model=MAPPER[MODEL])

prompt = st.text_input("Enter your regional bengali text:")
district = st.selectbox("Select District", [
    "Kishoreganj", "Narail", "Narsingdi", "Chittagong", "Rangpur", "Tangail"
])
if st.button("Generate Text"):
    if prompt != "" and district != "":
        ipa_transcription = model(f"<{district}> {prompt}",
                                  max_length=512)[0]["generated_text"]
        st.write(f"IPA Transcription:\n{ipa_transcription}")