Update app.py

app.py

@@ -3,13 +3,14 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Configurar la caché para evitar descargas repetitivas
+# Configurar caché y gestión de memoria
 os.environ["TRANSFORMERS_CACHE"] = "/root/.cache/huggingface/"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
 # Nombre del modelo
 model_name = "BSC-LT/ALIA-40b"
 
-# Intentar cargar el modelo desde la caché sin volver a descargarlo
+# Cargar modelo desde caché si es posible
 try:
     tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=os.getenv("TRANSFORMERS_CACHE"), local_files_only=True)
     model = AutoModelForCausalLM.from_pretrained(
@@ -17,7 +18,8 @@ try:
         cache_dir=os.getenv("TRANSFORMERS_CACHE"),
         local_files_only=True,
         device_map="auto",
-        torch_dtype=torch.bfloat16  # Precisión mixta para optimización de memoria
+        offload_folder="offload_cache",
+        torch_dtype=torch.bfloat16
     )
     print("Modelo cargado desde caché.")
 except Exception as e:
@@ -27,41 +29,31 @@ except Exception as e:
         model_name,
         cache_dir=os.getenv("TRANSFORMERS_CACHE"),
         device_map="auto",
-        torch_dtype=torch.bfloat16  # Precisión mixta para optimización de memoria
+        offload_folder="offload_cache",
+        torch_dtype=torch.bfloat16
     )
-
-    # Guardar el modelo localmente para futuras ejecuciones más rápidas
-    local_path = "/root/model_storage/"
-    tokenizer.save_pretrained(local_path)
-    model.save_pretrained(local_path)
+    tokenizer.save_pretrained("/root/model_storage/")
+    model.save_pretrained("/root/model_storage/")
     print("Modelo guardado en caché para futuras cargas.")
 
-# Verificar si la GPU está disponible y mover el modelo a la GPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+# Mostrar en qué dispositivo está el modelo
 print(f"Modelo cargado en: {next(model.parameters()).device}")
 
 def generar_texto(entrada):
-    # Liberar caché de la GPU antes de la inferencia para evitar problemas de memoria
-    torch.cuda.empty_cache()
+    torch.cuda.empty_cache()  # Liberar caché antes de inferencia
 
-    # Tokenizar la entrada y mover a GPU
-    input_ids = tokenizer(entrada, return_tensors="pt").input_ids.to(device)
+    input_ids = tokenizer(entrada, return_tensors="pt").input_ids.to("cuda")
 
-    # Generar texto con parámetros optimizados
     output = model.generate(
         input_ids,
-        max_length=100,        # Limitar longitud para ahorrar memoria
-        temperature=0.7,       # Controla la aleatoriedad
-        top_p=0.9,             # Reduce la búsqueda del modelo
-        num_return_sequences=1,# Genera una única respuesta
-        do_sample=True,        # Sampling para rapidez
-        use_cache=True         # Optimiza reutilizando cálculos previos
+        max_length=50,
+        temperature=0.7,
+        top_p=0.9,
+        num_return_sequences=1,
+        do_sample=True
     )
 
-    # Decodificar y retornar el texto generado
-    texto_generado = tokenizer.decode(output[0], skip_special_tokens=True)
-    return texto_generado
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
 # Crear la interfaz de Gradio
 interfaz = gr.Interface(