Update app.py

app.py

@@ -8,22 +8,22 @@ from threading import Thread
 huggingface_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
 login(huggingface_token)
 
-# Cambiar a la versión Meta Llama 3.1 1B
-model_id = "meta-llama/Llama-3.2-1B"
-
 # 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
 
-# Definir longitud máxima de tokens
 MAX_INPUT_TOKEN_LENGTH = 4096
 
-def generate_response(input_text, temperature=0.2, max_new_tokens=100):
-    """Función de generación de texto con el modelo."""
+def generate_response(input_text, temperature=0.7, max_new_tokens=100):
     input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    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=30.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         input_ids=input_ids,
         streamer=streamer,
@@ -31,58 +31,60 @@ def generate_response(input_text, temperature=0.2, max_new_tokens=100):
         do_sample=True,
         top_k=50,
         top_p=0.9,
-        num_beams=3,  # Beam search con 3 opciones
         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()
+    try:
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+        t.join()  # Esperar a que el hilo termine
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-
-    # Devolver la respuesta completa
-    return "".join(outputs)
+        outputs = []
+        for text in streamer:
+            outputs.append(text)
+        if not outputs:
+            raise ValueError("No se generó ninguna respuesta.")
+        return "".join(outputs)
+    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")
     
-    # Subir 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)
-        query = "aspiring human resources specialist"
-        st.write("Archivo CSV cargado exitosamente:")
-        st.write(df.head())  # Mostrar las primeras filas del dataframe
-        job_titles = df['job_title'].tolist()
+        
+        if 'job_title' in df.columns:
+            job_titles = df['job_title'].tolist()
+            query = "aspiring human resources specialist"
+            
+            st.write("Archivo CSV cargado exitosamente:")
+            st.write(df.head())
 
-        # Prompt mejorado
-        initial_prompt = f"Here is a list of job titles: {job_titles}. Please only return the first job title from this list without any additional text."
-        st.write(f"Query: {query}")
-        st.write(f"Prompt inicial: {initial_prompt}")
+            initial_prompt = f"Here is a list of job titles: {job_titles}. Could you extract the first job title from this list and print it clearly?"
+            st.write(f"Query: {query}")
+            st.write(f"Prompt inicial: {initial_prompt}")
+
+            if st.button("Generar respuesta"):
+                with st.spinner("Generando respuesta..."):
+                    response = generate_response(initial_prompt, temperature=0.2)
+                    if response:
+                        st.write(f"Respuesta del modelo: {response}")
+                    else:
+                        st.warning("No se pudo generar una respuesta.")
 
-        # Generar la respuesta del modelo
-        if st.button("Generar respuesta"):
-            with st.spinner("Generando respuesta..."):
-                response = generate_response(initial_prompt, temperature=0.7)
-                if response:
-                    st.write(f"Respuesta del modelo: {response}")
-                else:
-                    st.warning("No se pudo generar una respuesta.")
-                
-                # 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()
                 elif st.button("Terminar"):
                     st.stop()
+        else:
+            st.error("La columna 'job_title' no se encuentra en el archivo CSV.")
 
 if __name__ == "__main__":
     main()
-