from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import streamlit as st
from huggingface_hub import login
import pandas as pd
from threading import Thread

# Token Secret de Hugging Face
huggingface_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
login(huggingface_token)

# Cambiar a la versión Meta Llama 3.1 3B
model_id = "meta-llama/Llama-3.2-1B"

# Cargar el tokenizador y el modelo
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
tokenizer.pad_token = tokenizer.eos_token

# Definir longitud máxima de tokens
MAX_INPUT_TOKEN_LENGTH = 4096

def generate_response(input_text, temperature=0.2, max_new_tokens=50):
    """Función de generación de texto con el modelo."""
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=temperature != 0,
        top_k=10,  # Limita las palabras más probables a 50
        top_p=0.9,  # Considera solo el 90% de la probabilidad acumulada
        temperature=temperature,
        eos_token_id=[tokenizer.eos_token_id]
    )

    # Generación de texto en un hilo separado
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)

    # Devolver la respuesta completa
    return "".join(outputs)


def main():
    st.title("Chat con Meta Llama 3.1 3B")
    
    # Paso 1: Subir el archivo CSV
    uploaded_file = st.file_uploader("Por favor, sube un archivo CSV para iniciar:", type=["csv"])
    
    if uploaded_file is not None:
        df = pd.read_csv(uploaded_file)
        query = "aspiring human resources specialist"
        st.write("Archivo CSV cargado exitosamente:")
        st.write(df.head())  # Mostrar las primeras filas del dataframe
        job_titles = df['job_title'].tolist()
        # Prompt inicial
        initial_prompt = f"I have a list of job titles: {job_titles}. Please give me the first job title from the list."

        st.write(f"Query: {query}")
        st.write(f"Prompt inicial: {initial_prompt}")

        # Generar la respuesta del modelo
        if st.button("Generar respuesta"):
            with st.spinner("Generando respuesta..."):
                response = generate_response(initial_prompt)  # Obtener la primera respuesta completa
                st.write(f"Respuesta del modelo: {response}")

                # Terminar la conversación
                st.success("La conversación ha terminado.")
                
                # Opción para reiniciar o finalizar
                if st.button("Iniciar nueva conversación"):
                    st.experimental_rerun()  # Reinicia la aplicación
                elif st.button("Terminar"):
                    st.stop()

if __name__ == "__main__":
    main()