Actualizando API para modelo fusionado

api.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from huggingface_hub import hf_hub_download
+from huggingface_hub import snapshot_download
 import os
 
 # ✅ Inicializar FastAPI
@@ -11,43 +11,35 @@ app = FastAPI()
 # ✅ Definir un directorio de caché seguro
 os.environ["HF_HOME"] = "/tmp/huggingface"
 
-# ✅ Nombre del modelo en Hugging Face Hub
-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
+# ✅ Nombre del modelo en Hugging Face Hub (modelo fusionado)
+HUGGING_FACE_REPO = "fcp2207/Fusion_modelo_Phi2"  # Ruta donde subiste el modelo fusionado
 
-# ✅ 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
-)
+# ✅ Descargar el modelo fusionado desde Hugging Face
+print("🔄 Descargando modelo fusionado...")
+model_path = snapshot_download(repo_id=HUGGING_FACE_REPO, cache_dir=os.environ["HF_HOME"])
 
-# ✅ Cargar el tokenizer
-tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", cache_dir=os.environ["HF_HOME"])
+# ✅ Cargar el tokenizer desde el modelo fusionado
+print("🔄 Cargando tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 # ✅ Cargar el modelo en modo optimizado para memoria
+print("🔄 Cargando modelo...")
 model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/phi-2",
-    cache_dir=os.environ["HF_HOME"],
+    model_path,
     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
 
-
 # ✅ Definir la estructura de la solicitud para la API
 class InputText(BaseModel):
     input_text: str
 
-
 @app.get("/")
 def home():
     """Endpoint de prueba para verificar que la API está activa"""
-    return {"message": "API de Chatbot con Phi-2 está en funcionamiento 🚀"}
-
+    return {"message": "API de Chatbot con Phi-2 fusionado está en funcionamiento 🚀"}
 
 @app.post("/predict/")
 def predict(request: InputText):
@@ -60,7 +52,6 @@ def predict(request: InputText):
     response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"response": response_text}
 
-
 # ✅ Ejecución en modo local (opcional, no necesario en Hugging Face)
 if __name__ == "__main__":
     import uvicorn