app.py

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

# Asegurar que el token de padding esté configurado
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token  # Asignar el token de padding al token de fin de oración

def generate_response(input_text, max_new_tokens=50, temperature=0.5):
    """
    Genera una respuesta usando el modelo de lenguaje con máscara de atención.
    """
    # Tokenizar la entrada y crear la máscara de atención
    inputs = tokenizer(
        input_text,
        return_tensors='pt',
        padding=True,
        truncation=True,
        max_length=512  # Ajustar según sea necesario
    )

    input_ids = inputs['input_ids'].to(model.device)
    attention_mask = inputs['attention_mask'].to(model.device)

    # Generar texto con la máscara de atención y el token de padding
    outputs = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        pad_token_id=tokenizer.pad_token_id  # Usar el token de padding configurado
    )

    # Decodificar la respuesta generada
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

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:
            query = "aspiring human resources specialist"
            job_titles = df['job_title'].tolist()

            # Definir el prompt con in-context learning
            initial_prompt = (
                "You are an AI assistant specialized in job title extraction and similarity calculation.\n"
                "Your task is to extract the first job title from the following list and calculate the cosine similarity with the given query.\n"
                f"List: {job_titles}\n"
                f"Query: '{query}'\n"
                "Extracted first job title: \n"
                "Cosine similarity score: "
            )


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