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=20):
    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=120.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=20,
        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()  # Asegura que la generación haya terminado

        outputs = []
        for text in streamer:
            outputs.append(text)
        if not outputs:
            raise ValueError("No se generó ninguna respuesta.")
        
        response = "".join(outputs).strip().split("\n")[0]
        return response
    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()

            # Definir el prompt con in-context learning
            initial_prompt = (
                "Here are some examples of job title extraction:\n"
                "Example 1:\n"
                "List: ['Data Scientist', 'Machine Learning Engineer', 'AI Researcher']\n"
                "First job title: 'Data Scientist'\n"
                "\n"
                "Example 2:\n"
                "List: ['Software Developer', 'Backend Engineer', 'Frontend Developer']\n"
                "First job title: 'Software Developer'\n"
                "\n"
                "Now, extract the first job title from the following list:\n"
                f"List: {job_titles}\n"
                "First job title:"
            )
            
            st.write("Prompt inicial con In-context Learning:")
            st.write(initial_prompt)

            if st.button("Generar respuesta"):
                with st.spinner("Generando respuesta..."):
                    response = generate_response(initial_prompt, temperature=0.5)
                    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()