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

def generate_response(input_text, temperature=0.7, 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,
        num_return_sequences=3,
        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)
        query = 'aspiring human resources specialist'
        value = 0.00
        if 'job_title' in df.columns:
            job_titles = df['job_title']

            # Definir el prompt con in-context learning
            initial_prompt = (
                "Step 1: Extract the first record from the dataframe df.\n"
                f" {df.iloc[0]['job_title']}\n"
                #f"List: {job_titles}\n"
                #"First job title: \n"
                #"\n"
                "Step 2: Calculate the cosine similarity score between the job_title of the extracted record {df.iloc[0]['job_title']} and the given {query} and assign it to {value}.\n"
                f"Query: '{query}'\n"
                "Cosine similarity score: \n"
                "Step 3: Print the value of the calculated cosine similarity"
                f"Result: {value}"
            )

            
            st.write("Prompt inicial con In-context Learning:\n")
            st.write(initial_prompt)
            st.write(query)

            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()