import streamlit as st
from transformers import pipeline
import string

st.set_page_config(page_title="RLM-Translator", page_icon="🍕", layout="wide", initial_sidebar_state="collapsed")

@st.cache_resource
def load_model():
    return pipeline(model="rudrashah/RLM-hinglish-translator")

pipe = load_model()

def check_sentence_end(sentence):
    last_char = sentence[-1]
    if last_char in string.punctuation:
        return sentence
    else:
        return sentence + "."

def model_rlm(sentence):
    text = "mere paas 100 rupaye hain"
    text = check_sentence_end(sentence)

    template = "Hinglish:\n{hi_en}\n\nEnglish:\n{en}"
    english = pipe(template.format(hi_en=text,en=""), max_length=250)
    english = english[0]['translation_text']
    english = english.replace("<bos>","").replace("<eos>","")
    english = english[len(template.format(hi_en=text,en="")):]
    return english.strip()

def translate_hinglish_to_english(text):
    translated_text = model_rlm(text)
    st.session_state['translated_text'] = translated_text
    return translated_text

st.title("RLM-Translator")
st.write("A simple Hinglish to English translator from Hugging Face by [Rudra Shah](https://huggingface.co/rudrashah/RLM-hinglish-translator).")

col1, col2 = st.columns(2)

with col1:
    hinglish_text = st.text_area("Hinglish", height=300)

with col2:
    if "translated_text" in st.session_state:
        st.text_area("English", value=st.session_state["translated_text"], height=300, key="")
    else:
        st.text_area("English", height=300, key="")

if st.button("Translate", use_container_width=True, type="primary"):
    if hinglish_text != "":
        translated_text = translate_hinglish_to_english(hinglish_text)
        st.session_state["translated_text"] = translated_text
        st.rerun()
    else:
        st.error("Please enter some hinglish text to translate.")