Create app.py

app.py

@@ -0,0 +1,176 @@
+import os
+import re
+import json
+import requests
+import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from google.cloud import storage
+from google.auth import exceptions
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.hf_api import HfApi, HfFolder, HfLoginManager
+from io import BytesIO
+from dotenv import load_dotenv
+import uvicorn
+
+load_dotenv()
+
+# Variables de entorno
+API_KEY = os.getenv("API_KEY")
+GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
+GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
+HF_API_TOKEN = os.getenv("HF_API_TOKEN")  # Token de Hugging Face
+
+# Inicialización del cliente de GCS
+try:
+    credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
+    storage_client = storage.Client.from_service_account_info(credentials_info)
+    bucket = storage_client.bucket(GCS_BUCKET_NAME)
+except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError) as e:
+    print(f"Error al cargar credenciales o bucket: {e}")
+    exit(1)
+
+# Inicialización de FastAPI
+app = FastAPI()
+
+# Inicio de sesión en Hugging Face
+try:
+    if not HF_API_TOKEN:
+        raise ValueError("El token de Hugging Face no está definido en las variables de entorno.")
+    HfApi().set_access_token(HF_API_TOKEN)
+    print("Inicio de sesión en Hugging Face exitoso.")
+except Exception as e:
+    print(f"Error al iniciar sesión en Hugging Face: {e}")
+    exit(1)
+
+
+class DownloadModelRequest(BaseModel):
+    model_name: str
+    pipeline_task: str
+    input_text: str
+
+
+class GCSStreamHandler:
+    def __init__(self, bucket_name):
+        self.bucket = storage_client.bucket(bucket_name)
+
+    def file_exists(self, blob_name):
+        return self.bucket.blob(blob_name).exists()
+
+    def stream_file_from_gcs(self, blob_name):
+        blob = self.bucket.blob(blob_name)
+        if not blob.exists():
+            raise HTTPException(status_code=404, detail=f"Archivo '{blob_name}' no encontrado en GCS.")
+        return blob.download_as_bytes()
+
+    def upload_file_to_gcs(self, blob_name, data_stream):
+        blob = self.bucket.blob(blob_name)
+        blob.upload_from_file(data_stream)
+        print(f"Archivo {blob_name} subido a GCS.")
+
+    def ensure_bucket_structure(self, model_prefix):
+        # Crea automáticamente la estructura en el bucket si no existe
+        required_files = ["config.json", "tokenizer.json"]
+        for filename in required_files:
+            blob_name = f"{model_prefix}/{filename}"
+            if not self.file_exists(blob_name):
+                print(f"Creando archivo ficticio: {blob_name}")
+                self.bucket.blob(blob_name).upload_from_string("{}", content_type="application/json")
+
+    def stream_model_files(self, model_prefix, model_patterns):
+        model_files = {}
+        for pattern in model_patterns:
+            blobs = list(self.bucket.list_blobs(prefix=f"{model_prefix}/"))
+            for blob in blobs:
+                if re.match(pattern, blob.name.split('/')[-1]):
+                    print(f"Archivo encontrado: {blob.name}")
+                    model_files[blob.name.split('/')[-1]] = BytesIO(blob.download_as_bytes())
+        return model_files
+
+
+@app.post("/predict/")
+async def predict(request: DownloadModelRequest):
+    try:
+        gcs_handler = GCSStreamHandler(GCS_BUCKET_NAME)
+
+        # Asegura la estructura del bucket
+        gcs_handler.ensure_bucket_structure(request.model_name)
+
+        # Define patrones para los archivos de modelos
+        model_patterns = [
+            r"pytorch_model-\d+-of-\d+",
+            r"model-\d+",
+            r"pytorch_model.bin",
+            r"model.safetensors"
+        ]
+
+        # Carga los archivos del modelo desde el bucket
+        model_files = gcs_handler.stream_model_files(request.model_name, model_patterns)
+
+        # Cargar configuración y modelo
+        config_stream = gcs_handler.stream_file_from_gcs(f"{request.model_name}/config.json")
+        tokenizer_stream = gcs_handler.stream_file_from_gcs(f"{request.model_name}/tokenizer.json")
+
+        model = AutoModelForCausalLM.from_pretrained(BytesIO(config_stream))
+        state_dict = {}
+
+        for filename, stream in model_files.items():
+            state_dict.update(torch.load(stream, map_location="cpu"))
+
+        model.load_state_dict(state_dict)
+        tokenizer = AutoTokenizer.from_pretrained(BytesIO(tokenizer_stream))
+
+        # Crear pipeline
+        pipeline_task = request.pipeline_task
+        if pipeline_task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering"]:
+            raise HTTPException(status_code=400, detail="Unsupported pipeline task")
+
+        pipeline_ = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
+        input_text = request.input_text
+        result = pipeline_(input_text)
+
+        return {"response": result}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error: {e}")
+
+
+@app.post("/upload/")
+async def upload_model_to_gcs(model_name: str):
+    """
+    Descarga un modelo desde Hugging Face y lo sube a GCS en streaming.
+    """
+    try:
+        gcs_handler = GCSStreamHandler(GCS_BUCKET_NAME)
+
+        # Archivos comunes de los modelos
+        file_patterns = [
+            "pytorch_model.bin",
+            "model.safetensors",
+            "config.json",
+            "tokenizer.json",
+        ]
+
+        # Agregar patrones para fragmentos de modelos
+        for i in range(1, 100):
+            file_patterns.append(f"pytorch_model-{i:05}-of-{100:05}")
+            file_patterns.append(f"model-{i:05}")
+
+        for filename in file_patterns:
+            url = f"https://huggingface.co/{model_name}/resolve/main/{filename}"
+            headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
+            try:
+                response = requests.get(url, headers=headers, stream=True)
+                if response.status_code == 200:
+                    blob_name = f"{model_name}/{filename}"
+                    blob = bucket.blob(blob_name)
+                    blob.upload_from_file(BytesIO(response.content))
+                    print(f"Archivo {filename} subido correctamente a GCS.")
+            except Exception as e:
+                print(f"Archivo {filename} no encontrado: {e}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error al subir modelo: {e}")
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)