from transformers import pipeline, AutoTokenizer
import streamlit as st
import os

# model_id = "meta-llama/Llama-3.1-8B"
model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
os.makedirs("/app/cache", exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir="/app/cache")
generator = pipeline(
    'text-generation',
    model=model_id,
    tokenizer=tokenizer,
    device_map="auto",
)

context = {
    "role":"system",
    "content":
    """
    Seu nome é IASYS - Sistema de Assistência Inteligente para Saúde.
Você é um chatbot prestativo e simpático, responde em Português, e está sendo desenvolvido por Mariana Lima e Aylla Matos, 
ambas da Faculdade de Petrolina (FACAPE), para que elas possam aprender sobre desenvolvimento de chatbots
na prática.
Você está sendo desenvolvido para auxiliar em um projeto de extensão acadêmica da Universidade de Pernambuco (UPE)
e da Faculdade de Petrolina (FACAPE) chamado "PET Saúde Digital", que visa auxiliar a saúde digital
no sistema de atenção primária de Petrolina.
Você deve demonstrar carisma e educação ao falar, mas não se gabar disso.
Lembre-se de falar como alguém que está falando de si mesmo, não de outra pessoa.
Fale mais sobre seu papel no projeto e sobre o seu entusiasmo em auxiliar a atenção primária de Petrolina.
    """
}

st.set_page_config(
    page_title = "IASYS",
    page_icon= "👩🏻‍🔬",
)
st.title('IASYS - Intelligent Assistant System for Health')

if 'messages' not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    st.chat_message(message['role']).markdown(message['content'])

prompt = st.chat_input()
system_prompt = context["content"]

if prompt:
    st.chat_message('user').markdown(prompt)
    st.session_state.messages.append({'role':'user', 'content': prompt})

    with st.spinner("IASYS is typing..."):
        full_prompt = tokenizer.apply_chat_template([context, {'role':'user', 'content': prompt}],
            tokenize=False,
            add_generation_prompt=True
        )
        response = generator(
            full_prompt,
            max_length=380,
            min_length=150,
            do_sample=True,
            truncation=True,
            temperature=0.7,
        )

        generated_text = response[0]['generated_text']
        print(generated_text)

    st.chat_message('assistant').markdown(generated_text)
    st.session_state.messages.append({'role':'assistant', 'content': generated_text})