Update app.py

app.py

@@ -2,78 +2,76 @@ 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_name = "meta-llama/Llama-3.2-1B-Instruct"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Asignar el eos_token como pad_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
 
-# Control deslizante para ajustar la temperatura
-temperature = st.slider("Ajustar temperatura", min_value=0.1, max_value=1.5, value=0.7, step=0.1)
-
-# Upload CSV file
-uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
+# Definir longitud máxima de tokens
+MAX_INPUT_TOKEN_LENGTH = 4096
 
-# Leer el archivo CSV si se ha subido
-if uploaded_file is not None:
-    df = pd.read_csv(uploaded_file)
-    st.write(df.head())  # Mostrar las primeras filas del dataframe
+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)
 
-    # Verificar si la columna 'job_title' está en el dataframe
-    if 'job_title' in df.columns:
-        job_titles = df['job_title'].tolist()
-    else:
-        st.error("La columna 'job_title' no se encuentra en el archivo CSV.")
-        job_titles = []  # Asignar una lista vacía si la columna no existe
+    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]
+    )
 
-else:
-    st.warning("Por favor, sube un archivo CSV.")
-    job_titles = []  # Asignar una lista vacía si no se ha subido un archivo
+    # Generación de texto en un hilo separado
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
 
-# Definir la consulta
-query = "aspiring human resources specialist"
-st.write("Query:", query)
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
 
-# Texto de entrada para la generación
-input_text = (
-    f"You are an AI assistant. You have a list of job titles and a search query.\n"
-    f"Your task is to rank these job titles by their semantic similarity to the given query. "
-    f"Please provide the ranking from most relevant to least relevant. "
-    f"Do not calculate cosine similarity; instead, focus on understanding the semantic relevance of each job title to the query.\n"
-    f"\n"
-    f"Format your response like this:\n"
-    f"1. [Most Relevant Job Title]\n"
-    f"2. [Second Most Relevant Job Title]\n"
-    f"...\n"
-    f"N. [Least Relevant Job Title]\n"
-    f"\n"
-    f"Query: \"{query}\"\n"
-    f"Job Titles: {job_titles}\n"
-)
+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
 
-st.write("Texto de entrada para la generación:", input_text)
+        # Prompt inicial
+        initial_prompt = "dame el nombre de un animal"
+        st.write(f"Prompt inicial: {initial_prompt}")
 
-# Codificar el texto de entrada
-input_ids = tokenizer.encode(input_text, return_tensors='pt')
+        # 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}")
 
-# Generar texto utilizando la temperatura ajustada
-outputs = model.generate(
-    input_ids,
-    max_length=2000,
-    temperature=temperature,  # Usar la temperatura ajustada
-    top_k=50,
-    top_p=0.9,
-    do_sample=True
-)
+                # 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()
 
-# Decodificar el texto generado
-generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-st.write("Texto generado:", generated_text)
+if __name__ == "__main__":
+    main()