Optimización de memoria en Hugging Face Spaces

api.py

@@ -15,18 +15,25 @@ os.environ["HF_HOME"] = "/tmp/huggingface"
 HUGGING_FACE_REPO = "fcp2207/Phi-2"  # Reemplaza con tu usuario y nombre correcto del modelo en Hugging Face
 MODEL_FILENAME = "phi2_finetuned.pth"  # Nombre del archivo en Hugging Face
 
-# ✅ Descargar el modelo desde Hugging Face (usando la caché segura)
+# ✅ Descargar el modelo desde Hugging Face con caché segura
 model_path = hf_hub_download(
     repo_id=HUGGING_FACE_REPO,
     filename=MODEL_FILENAME,
     cache_dir=os.environ["HF_HOME"]  # Directorio seguro en Hugging Face Spaces
 )
 
-# ✅ Cargar el tokenizer y el modelo base desde Hugging Face
+# ✅ Cargar el tokenizer
 tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", cache_dir=os.environ["HF_HOME"])
-model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", cache_dir=os.environ["HF_HOME"])
 
-# ✅ Cargar los pesos del modelo fine-tuned
+# ✅ Cargar el modelo en modo optimizado para memoria
+model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/phi-2",
+    cache_dir=os.environ["HF_HOME"],
+    torch_dtype=torch.float16,  # Reduce el tamaño del modelo
+    device_map="auto"  # Optimiza la carga en CPU/GPU automáticamente
+)
+
+# ✅ Cargar los pesos del modelo entrenado
 model.load_state_dict(torch.load(model_path, map_location="cpu"))
 model.eval()  # Poner el modelo en modo inferencia
 
@@ -68,3 +75,4 @@ if __name__ == "__main__":
 
 
 
+