import streamlit as st
import maritalk

def init_model():
    model = maritalk.MariTalk(
        key="100479706496023140829_534ece7d352abf00",
        model="sabia-3"  # No momento, suportamos os modelos sabia-3, sabia-2-medium e sabia-2-small
    )
    return model

def stream_response(model, message, max_tokens=200):
    for response in model.generate(message, max_tokens=max_tokens, stream=True): 
        if 'text' in response.keys():
            yield response['text']
        
# main()
model = init_model()
with st.container():
    st.write('\n\n')
    st.title('LLM-LANAChat-Brasil\n\n')
    message = st.text_input("Chat", placeholder="E ai?")
    if message:
        max_tokens = 200
        st.write_stream(stream_response(model, message, max_tokens))