Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -6,48 +5,57 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # Especifica el nombre del modelo
 model_name = "BSC-LT/ALIA-40b"
 
-# Cargar el tokenizador y el modelo
+# Verificar si hay GPU disponible y configurarla
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Usando dispositivo: {device}")
+
+# Cargar el tokenizador y el modelo con optimizaciones
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    device_map="auto",
-    torch_dtype=torch.bfloat16  # Reduce uso de memoria a la mitad
+    device_map="auto",  # Distribuye el modelo automáticamente entre las GPUs disponibles
+    torch_dtype=torch.bfloat16,  # Reduce uso de memoria
+    offload_folder="offload_cache"  # Guarda partes del modelo en disco si es necesario
 )
 
-# Verificar en qué dispositivo está el modelo
-print(f"Modelo cargado en: {model.device}")
-
+# Mover modelo al dispositivo seleccionado
+model.to(device)
 
+# Verificar en qué dispositivo está cargado el modelo
+print(f"Modelo cargado en: {next(model.parameters()).device}")
 
 def generar_texto(entrada):
-    # Tokenizar la entrada
-    input_ids = tokenizer.encode(entrada, return_tensors="pt").input_ids.to("cuda")
-    
-    # Generar texto con el modelo
+    # Liberar memoria de la GPU antes de la inferencia
+    torch.cuda.empty_cache()
+
+    # Tokenizar la entrada y mover a GPU
+    input_ids = tokenizer(entrada, return_tensors="pt").input_ids.to(device)
+
+    # Configurar la generación con parámetros optimizados
     output = model.generate(
         input_ids,
-        max_length=100,         # Reduce la longitud de salida
-        temperature=0.7,        # Controla la aleatoriedad
-        top_p=0.9,              # Reduce la búsqueda del modelo
-        num_return_sequences=1, # Evita múltiples respuestas simultáneas
-        do_sample=True          # Sampling para rapidez
+        max_length=50,          # Ajuste del tamaño de salida para optimizar la memoria
+        temperature=0.7,         # Controla la aleatoriedad
+        top_p=0.9,               # Reduce la búsqueda del modelo
+        num_return_sequences=1,  # Generar solo una respuesta
+        do_sample=True,          # Sampling para rapidez
+        use_cache=True           # Optimiza reutilizando cálculos previos
     )
 
-    
     # Decodificar y retornar el texto generado
     texto_generado = tokenizer.decode(output[0], skip_special_tokens=True)
     return texto_generado
 
-# Crear la interfaz de Gradio usando la sintaxis actualizada
+# Crear la interfaz de Gradio con entradas interactivas
 interfaz = gr.Interface(
     fn=generar_texto,
-    inputs=gr.Textbox(lines=2, placeholder="Escribe tu prompt aquí..."),
-    outputs=gr.Textbox(),
+    inputs=gr.Textbox(lines=2, placeholder="Escribe tu prompt aquí...", interactive=True),
+    outputs=gr.Textbox(interactive=True),
     title="Generador de Texto con ALIA-40b",
     description="Este modelo genera texto utilizando ALIA-40b, un modelo LLM entrenado por BSC-LT."
 )
 
-
-
+# Lanzar la interfaz con opción de compartir
 if __name__ == "__main__":
-    interfaz.launch()
+    interfaz.launch(share=True, server_port=7860)