from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from accelerate import init_empty_weights, load_checkpoint_and_dispatch, dispatch_model, infer_auto_device_map
import streamlit as st
from huggingface_hub import login
import pandas as pd

# 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

# Cargar el modelo con disk_offload
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(model_id)

device_map = infer_auto_device_map(model, max_memory={"disk": "2GiB"}, no_split_module_classes=["LlamaDecoderLayer"])
model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, offload_folder="offload_dir")

MAX_INPUT_TOKEN_LENGTH = 10000


def generate_response(input_text, temperature=0.7, max_new_tokens=20):
    input_ids = tokenizer.encode(input_text, return_tensors='pt').to("cpu")  # Usar 'cpu' para mantener la compatibilidad

    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,
        num_return_sequences=3,
        eos_token_id=tokenizer.eos_token_id
    )

    try:
        outputs = model.generate(**generate_kwargs)

        response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
        return response.split("\n")[0]
    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)
        query = 'aspiring human resources specialist'
        
        if 'job_title' in df.columns:
            job_titles = df['job_title'].tolist()

            # Definir el prompt con in-context learning
            initial_prompt = (
                f"Extract the first record from the dataframe df.\n"
                f"First job title: '{df.iloc[0]['job_title']}'\n"
                f"Calculate the cosine similarity between this job title and the query: '{query}'.\n"
                "Print the cosine similarity score."
            )

            st.write("Prompt inicial con In-context Learning:\n")
            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()