import streamlit as st
from langchain.llms import HuggingFaceHub
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
  

# Answer:
# [
# 0:{
# "generated_text":"I like rice I like rice I like rice I like ri"
# }
# ]

#Function to return the response
def load_answer(question):
    tokenizer = AutoTokenizer.from_pretrained("unicamp-dl/translation-pt-en-t5")
    model = AutoModelForSeq2SeqLM.from_pretrained("unicamp-dl/translation-pt-en-t5")
    pten_pipeline = pipeline('text2text-generation', model=model, tokenizer=tokenizer)
    return pten_pipeline(question)[0]["generated_text"]

#App UI starts here
st.set_page_config(page_title="Lanza Chatbot", page_icon=":robot:")
st.header("Lanza Chatbot")

#Gets the user input
def get_text():
    input_text = st.text_input("You: ", key="input")
    return input_text

user_input=get_text()
response = load_answer(user_input)

submit = st.button('Generate')

if submit:

    st.subheader("Answer:")
    st.write(response)