app.py

from transformers import AutoModelForCausalLM, AutoTokenizer
import streamlit as st
from huggingface_hub import login
import pandas as pd
from threading import Thread

# Token Secret of Hugging Face
huggingface_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
login(huggingface_token)


# Cargar el modelo y el tokenizer
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
tokenizer.pad_token = tokenizer.eos_token

# Definir longitud máxima de tokens
MAX_INPUT_TOKEN_LENGTH = 4096

def generate_response(input_text, temperature=0.7, max_new_tokens=100):
    """Función de generación de texto con el modelo."""
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.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=temperature != 0,
        temperature=temperature,
        eos_token_id=[tokenizer.eos_token_id]
    )

    # Generación de texto en un hilo separado
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)

def main():
    st.title("Chat con Meta Llama 3.1 8B")
    
    # Paso 1: Subir el archivo CSV
    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)
        st.write("Archivo CSV cargado exitosamente:")
        st.write(df.head())  # Mostrar las primeras filas del dataframe

        # Prompt inicial
        initial_prompt = "dame el nombre de un animal"
        st.write(f"Prompt inicial: {initial_prompt}")

        # Generar la respuesta del modelo
        if st.button("Generar respuesta"):
            with st.spinner("Generando respuesta..."):
                response = list(generate_response(initial_prompt))[0]  # Obtener la primera respuesta completa
                st.write(f"Respuesta del modelo: {response}")

                # Terminar la conversación
                st.success("La conversación ha terminado.")
                
                # Opción para reiniciar o finalizar
                if st.button("Iniciar nueva conversación"):
                    st.experimental_rerun()  # Reinicia la aplicación
                elif st.button("Terminar"):
                    st.stop()

if __name__ == "__main__":
    main()