Update app.py

app.py

@@ -4,13 +4,14 @@ import json
 import requests
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from google.auth import exceptions
 from google.cloud import storage
+from google.auth import exceptions
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from io import BytesIO
 from dotenv import load_dotenv
 import uvicorn
 
+# Carga de variables de entorno
 load_dotenv()
 
 API_KEY = os.getenv("API_KEY")
@@ -18,35 +19,28 @@ 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")
 
-def sanitize_bucket_name(bucket_name):
-    """Corrige un nombre de bucket inválido para ajustarse a las reglas de Google Cloud Storage."""
-    bucket_name = bucket_name.lower()
-    bucket_name = re.sub(r"[^a-z0-9-\.]", "-", bucket_name)
-    bucket_name = bucket_name.strip("-.")
-    if len(bucket_name) > 63:
-        bucket_name = bucket_name[:63]
-    if not re.match(r"^[a-z0-9]", bucket_name):
-        bucket_name = "a" + bucket_name
-    if not re.match(r"[a-z0-9]$", bucket_name):
-        bucket_name = bucket_name + "a"
-    return bucket_name
-
 def validate_bucket_name(bucket_name):
-    """Valida si el nombre de bucket cumple con las reglas de Google Cloud Storage."""
-    if not re.match(r"^[a-z0-9][a-z0-9\-\.]*[a-z0-9]$", bucket_name):
-        raise ValueError(f"Nombre de bucket inválido: '{bucket_name}'. Debe cumplir con las reglas de GCS.")
+    """Valida que el nombre del bucket cumpla con las restricciones de Google Cloud."""
+    if not re.match(r"^[a-z0-9][a-z0-9\-]*[a-z0-9]$", bucket_name):
+        raise ValueError(f"Invalid bucket name '{bucket_name}'. Must start and end with a letter or number.")
     return bucket_name
 
+def validate_huggingface_repo_name(repo_name):
+    """Valida que el nombre del repositorio cumpla con las restricciones de Hugging Face."""
+    if not re.match(r"^[a-zA-Z0-9_.-]+$", repo_name):
+        raise ValueError(f"Invalid repository name '{repo_name}'. Must use alphanumeric characters, '-', '_', or '.'.")
+    if repo_name.startswith(('-', '.')) or repo_name.endswith(('-', '.')) or '..' in repo_name:
+        raise ValueError(f"Invalid repository name '{repo_name}'. Cannot start or end with '-' or '.', or contain '..'.")
+    if len(repo_name) > 96:
+        raise ValueError(f"Repository name '{repo_name}' exceeds max length of 96 characters.")
+    return repo_name
+
+# Validar y configurar cliente de GCS
 try:
-    # Sanitizar y validar el nombre del bucket
-    GCS_BUCKET_NAME = sanitize_bucket_name(GCS_BUCKET_NAME)
     GCS_BUCKET_NAME = validate_bucket_name(GCS_BUCKET_NAME)
-
-    # Cargar credenciales de Google Cloud Storage
     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, ValueError) as e:
     print(f"Error al cargar credenciales o bucket: {e}")
     exit(1)
@@ -58,49 +52,34 @@ class DownloadModelRequest(BaseModel):
     pipeline_task: str
     input_text: str
 
-class GCSStreamHandler:
+class GCSHandler:
     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):
+    def upload_file(self, blob_name, file_stream):
         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()
+        blob.upload_from_file(file_stream)
 
-    def upload_file_to_gcs(self, blob_name, data_stream):
+    def download_file(self, blob_name):
         blob = self.bucket.blob(blob_name)
-        blob.upload_from_file(data_stream)
-
-    def ensure_bucket_structure(self, model_prefix):
-        required_files = ["config.json", "tokenizer.json"]
-        for filename in required_files:
-            blob_name = f"{model_prefix}/{filename}"
-            if not self.file_exists(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]):
-                    model_files[blob.name.split('/')[-1]] = BytesIO(blob.download_as_bytes())
-        return model_files
+        if not blob.exists():
+            raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
+        return BytesIO(blob.download_as_bytes())
 
 def download_model_from_huggingface(model_name):
+    """Descarga un modelo desde Hugging Face y lo sube a GCS."""
+    model_name = validate_huggingface_repo_name(model_name)
     file_patterns = [
         "pytorch_model.bin",
-        "model.safetensors",
         "config.json",
         "tokenizer.json",
+        "model.safetensors",
     ]
     for i in range(1, 100):
-        file_patterns.append(f"pytorch_model-{i:05}-of-{100:05}")
-        file_patterns.append(f"model-{i:05}")
+        file_patterns.extend([f"pytorch_model-{i:05}-of-00001", 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}"}
@@ -108,37 +87,34 @@ def download_model_from_huggingface(model_name):
             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))
-        except Exception:
-            pass
+                bucket.blob(blob_name).upload_from_file(BytesIO(response.content))
+        except Exception as e:
+            print(f"Error downloading {filename} from Hugging Face: {e}")
 
 @app.post("/predict/")
 async def predict(request: DownloadModelRequest):
     try:
-        gcs_handler = GCSStreamHandler(GCS_BUCKET_NAME)
+        gcs_handler = GCSHandler(GCS_BUCKET_NAME)
         model_prefix = request.model_name
-        model_patterns = [
-            r"pytorch_model-\d+-of-\d+",
-            r"model-\d+",
-            r"pytorch_model.bin",
-            r"model.safetensors",
+        model_files = [
+            "pytorch_model.bin",
+            "config.json",
+            "tokenizer.json",
+            "model.safetensors",
         ]
-        if not any(
-            gcs_handler.file_exists(f"{model_prefix}/{pattern}") for pattern in model_patterns
-        ):
+        for i in range(1, 100):
+            model_files.extend([f"pytorch_model-{i:05}-of-00001", f"model-{i:05}"])
+        if not any(gcs_handler.file_exists(f"{model_prefix}/{file}") for file in model_files):
             download_model_from_huggingface(model_prefix)
-        model_files = gcs_handler.stream_model_files(model_prefix, model_patterns)
-        config_stream = gcs_handler.stream_file_from_gcs(f"{model_prefix}/config.json")
-        tokenizer_stream = gcs_handler.stream_file_from_gcs(f"{model_prefix}/tokenizer.json")
-        model = AutoModelForCausalLM.from_pretrained(BytesIO(config_stream))
-        tokenizer = AutoTokenizer.from_pretrained(BytesIO(tokenizer_stream))
-        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="Tarea no soportada")
-        pipeline_ = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
-        input_text = request.input_text
-        result = pipeline_(input_text)
+        model_files_streams = {file: gcs_handler.download_file(f"{model_prefix}/{file}") for file in model_files if gcs_handler.file_exists(f"{model_prefix}/{file}")}
+        config_stream = model_files_streams.get("config.json")
+        tokenizer_stream = model_files_streams.get("tokenizer.json")
+        if not config_stream or not tokenizer_stream:
+            raise HTTPException(status_code=500, detail="Required model files missing.")
+        model = AutoModelForCausalLM.from_pretrained(config_stream)
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
+        pipeline_ = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
+        result = pipeline_(request.input_text)
         return {"response": result}
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error: {e}")