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)

# Cargar el tokenizador y el modelo
model_id = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
tokenizer.pad_token = tokenizer.eos_token

MAX_INPUT_TOKEN_LENGTH = 10000

def generate_response(input_text, temperature=0.5, max_new_tokens=100):
    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)

    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        st.warning(f"Se recortó la entrada porque excedió el límite de {MAX_INPUT_TOKEN_LENGTH} tokens.")

    streamer = TextIteratorStreamer(tokenizer, timeout=60.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=True,
        top_k=50,
        top_p=0.9,
        temperature=temperature,
        eos_token_id=[tokenizer.eos_token_id]
    )

    try:
        t = Thread(target=model.generate, kwargs=generate_kwargs)
        t.start()
        t.join()  # Esperar a que el hilo termine

        outputs = []
        for text in streamer:
            outputs.append(text)
        if not outputs:
            raise ValueError("No se generó ninguna respuesta.")
        return "".join(outputs)
    except Exception as e:
        st.error(f"Error durante la generación: {e}")
        return "Error en la generación de texto."

def main():
    st.title("Chat con Meta Llama 3.2 1B")
    
    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)
        
        if 'job_title' in df.columns:
            job_titles = df['job_title'].tolist()
            query = "aspiring human resources specialist"
            
            st.write("Archivo CSV cargado exitosamente:")
            st.write(df.head())

            initial_prompt = f"I have a list of job titles: {job_titles}. Please extract and return only the first job title from this list without repeating."
            st.write(f"Query: {query}")
            st.write(f"Prompt inicial: {initial_prompt}")

            if st.button("Generar respuesta"):
                with st.spinner("Generando respuesta..."):
                    response = generate_response(initial_prompt, temperature=0.2)
                    if response:
                        st.write(f"Respuesta del modelo: {response}")
                    else:
                        st.warning("No se pudo generar una respuesta.")

                st.success("La conversación ha terminado.")
                
                if st.button("Iniciar nueva conversación"):
                    st.experimental_rerun()
                elif st.button("Terminar"):
                    st.stop()
        else:
            st.error("La columna 'job_title' no se encuentra en el archivo CSV.")

if __name__ == "__main__":
    main()