src/streamlit_app.py

import streamlit as st
from transformers import pipeline
import os

# ---------------- CONFIG ----------------
st.set_page_config(page_title="Kriolu AI Hub", layout="wide")
token = os.environ.get("token")

# ---------------- CACHE ----------------
@st.cache_resource
def load_pipeline(task, model_path, **kwargs):
    return pipeline(task, model=model_path, tokenizer=model_path, token=token, **kwargs)

# ---------------- DECODER ----------------
def instantiate_gpt2(model_name, max_length_, num_return_sequences, text):
    model_path = f'Iscte-Sintra/{model_name}'
    pipe = load_pipeline("text-generation", model_path)

    return pipe(
        text,
        max_new_tokens=max_length_,
        num_return_sequences=num_return_sequences,
        do_sample=True,
        top_p=0.95,
        top_k=50
    )

# ---------------- CHATBOT ----------------
def build_chatbot_page(model_name):
    st.title(f"🤖 {model_name}: Chatbot")
    
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Display previous messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Input new message
    if prompt := st.chat_input("Diz algo..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            with st.spinner("A pensar..."):
                try:
                    # --- PHASE 1: Load Pipeline & Tokenizer ---
                    pipe = load_pipeline("text-generation", f"Iscte-Sintra/{model_name}")
                    if pipe is None:
                        st.error("Falha ao inicializar o pipeline.")
                        return 
                    
                    tokenizer = pipe.tokenizer

                    # Inject ChatML template if missing
                    if getattr(tokenizer, "chat_template", None) is None:
                        tokenizer.chat_template = (
                            "{% for message in messages %}"
                            "{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n'}}"
                            "{% endfor %}"
                            "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}"
                        )

                    # --- PHASE 2: Format History ---
                    conversation = [{"role": "system", "content": "You are a helpful assistant."}] + st.session_state.messages
                    
                    prompt_str = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
                    
                    # Defensively ensure prompt_str is a string
                    if isinstance(prompt_str, list):
                        prompt_str = prompt_str[0]

                    # --- PHASE 3: Safe Token IDs ---
                    # Prevent the 'list' error by ensuring pad_id is strictly an integer
                    pad_id = tokenizer.eos_token_id
                    if isinstance(pad_id, list):
                        pad_id = pad_id[0] # Take the first token if the model has multiple
                    if pad_id is None:
                        pad_id = 50256

                    # --- PHASE 4: Generate ---
                    response = pipe(
                        str(prompt_str), # Force string conversion
                        max_new_tokens=150,
                        do_sample=True,
                        temperature=0.7,
                        pad_token_id=pad_id,
                        return_full_text=False
                    )
                    
                    # Defensively extract the generated text to handle different list structures
                    if isinstance(response, list) and len(response) > 0:
                        if isinstance(response[0], list): 
                            generated_text = response[0][0].get('generated_text', '')
                        else:
                            generated_text = response[0].get('generated_text', '')
                    else:
                        generated_text = str(response)
                    
                    generated_text = generated_text.strip()
                    
                    st.markdown(generated_text)
                    st.session_state.messages.append({"role": "assistant", "content": generated_text})
                    
                except Exception as e:
                    st.error(f"Erro na geração: {e}")
                    
# ---------------- ENCODER ----------------
def instantiate_encoder(model_name, top_k, text):
    pipe = load_pipeline("fill-mask", f"Iscte-Sintra/{model_name}")
    return pipe(text, top_k=top_k)

# ---------------- TRANSLATION ----------------
def instantiate_translation_model(model_name, text, src_lg, tgt_lg):
    model_path = f'Iscte-Sintra/{model_name}'

    # ---- NLLB ----
    if "nllb" in model_name:
        pipe = pipeline(
            "translation",
            model=model_path,
            tokenizer=model_path,
            token=token,
            src_lang=src_lg,
            tgt_lang=tgt_lg
        )
        return pipe(text)[0]["translation_text"]

    # ---- M2M100 ----
    elif "m2m100" in model_name:
        pipe = load_pipeline("translation", model_path)

        # 1. Definimos a língua de origem
        # Em modelos customizados, as vezes o src_lang precisa ser o token completo
        pipe.tokenizer.src_lang = src_lg 
        
        # 2. Pegamos o ID numérico do token de destino (ex: __pt__)
        # Usamos convert_tokens_to_ids porque ele ignora a lógica interna de busca de idiomas
        tgt_lang_id = pipe.tokenizer.convert_tokens_to_ids(tgt_lg)
        
        if tgt_lang_id == pipe.tokenizer.unk_token_id:
            st.error(f"Erro: O token {tgt_lg} não foi encontrado no vocabulário do modelo!")
            return None

        # 3. Executamos a tradução forçando o ID de início de frase (BOS)
        result = pipe(
            text,
            forced_bos_token_id=tgt_lang_id
        )
        return result[0]["translation_text"]
    # ---- MBART ----
    else:
        pipe = pipeline(
            "translation",
            model=model_path,
            tokenizer=model_path,
            token=token,
            src_lang=src_lg,
            tgt_lang=tgt_lg
        )
        return pipe(text)[0]["translation_text"]

# ---------------- UI: TRANSLATION ----------------
def build_translation_page(model_name):
    st.title(f"🌍 {model_name}: Tradução")

    if "nllb" in model_name:
        lang_map = {
            "Português": "por_Latn",
            "Kabuverdianu": "kea_Latn"
        }

    elif "m2m100" in model_name:
        lang_map = {
            "Português": "__pt__",
            "Kabuverdianu": "__en__"  # Proxying kea as __en__
        }

    else:  # mBART
        lang_map = {
            "Português": "pt_XX",
            "Kabuverdianu": "en_XX"
        }

    col1, col2 = st.columns(2)
    with col1:
        src_label = st.selectbox("Língua de Origem", list(lang_map.keys()))
    with col2:
        tgt_label = st.selectbox("Língua de Destino", list(lang_map.keys()))

    text = st.text_area("Texto de entrada", "Katxór sta trás di pórta.", height=100)

    if st.button("Traduzir"):
        if not text.strip():
            st.warning("Introduza texto!")
            return

        with st.spinner("A traduzir..."):
            try:
                result = instantiate_translation_model(
                    model_name,
                    text,
                    lang_map[src_label],
                    lang_map[tgt_label]
                )
                st.success("Resultado:")
                st.write(result)
            except Exception as e:
                st.error(f"Erro: {e}")

# ---------------- UI: DECODER ----------------
def build_decoder_page(model_name):
    st.title(f"✍️ {model_name}: Geração de Texto")
    max_length = st.sidebar.slider("Máximo de Tokens", 10, 200, 50)
    num_seq = st.sidebar.number_input("Sequências", 1, 5, 1)
    text = st.text_area("Prompt", "Katxór sta trás di pórta.")

    if st.button("Gerar"):
        with st.spinner("A processar..."):
            try:
                results = instantiate_gpt2(model_name, max_length, num_seq, text)
                for res in results:
                    st.info(res["generated_text"])
            except Exception as e:
                st.error(f"Erro: {e}")

# ---------------- UI: ENCODER ----------------
def build_encoder_page(model_name):
    st.title(f"🔍 {model_name}: Fill-Mask")
    top_k = st.sidebar.slider("Top K sugestões", 1, 5, 3)

    mask_token = "<mask>" if "RoBERTa" in model_name else "[MASK]"
    st.write(f"Use o token **{mask_token}** para a palavra em falta.")

    input_text = st.text_input("Frase", f"Katxór sta trás di {mask_token}.")

    if st.button("Prever"):
        try:
            results = instantiate_encoder(model_name, top_k, input_text)
            for res in results:
                st.write(f"✅ **{res['token_str']}** ({res['score']:.2%})")
        except Exception:
            st.error(f"Certifique-se que usou o token {mask_token}")

# ---------------- MAIN ----------------
model_dict = {
    "RoBERTa-Kriolu": "Encoder",
    "GPT2_v1.18": "Decoder",
    "LLM-kea-v1.0": "Decoder",
    "Modelo-Traducao-kea-ptpt-v1.0": "Encoder-Decoder",
    "nllb-v1.0": "Encoder-Decoder",
    "m2m100-v1.0": "Encoder-Decoder",
    "mbart-v2.0": "Encoder-Decoder",
    "m2m100-v2.0": "Encoder-Decoder", 
    "mbart-v2.1": "Encoder-Decoder",
    "nllb-v2.0": "Encoder-Decoder",
    "m2m100-v2.1": "Encoder-Decoder",
    "mbart-v2.3": "Encoder-Decoder",
    "m2m100-v2.1": "Encoder-Decoder",
    "nllb-kea-v2.1": "Encoder-Decoder",
    "portuguese-experiment": "Chatbot",
}

selected_model = st.sidebar.selectbox("Escolha o Modelo", list(model_dict.keys()))
arch = model_dict[selected_model]

if arch == "Chatbot":
    build_chatbot_page(selected_model) # New call
elif arch == "Encoder":
    build_encoder_page(selected_model)
elif arch == "Encoder-Decoder":
    build_translation_page(selected_model)
else:
    build_decoder_page(selected_model)