import streamlit as st
import transformers
from transformers import pipeline
import indicnlp
from indicnlp.transliterate.unicode_transliterate import UnicodeIndicTransliterator

model_checkpoint = './MKB_NIOS_hindi_sanskrit_5'
translator = pipeline('translation', model=model_checkpoint, decoder_start_token='<2sa>')

st.title('NLTM')

st.header('Hindi-Sanskrit Translation Model')
query = st.text_input("Enter the Input Sentence", "")
submit = st.button('Translate') 
input_sentence = query + ' </s>' + ' <2hi>'

if submit:
  st.subheader('Sanskrit Sentence')
  
  with st.spinner(text='This may take a moment...'):
    output_sentence = translator(input_sentence,max_length=128)
    
  output_sentence = output_sentence[0]["translation_text"]
  output_sentence = output_sentence.replace('<2sa>','')
    
  st.write(output_sentence)