Update app.py

app.py

@@ -1,24 +1,19 @@
-from huggingface_hub import HfApi
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel
 import requests
 import boto3
 from dotenv import load_dotenv
 import os
 import uvicorn
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoConfig, TextIteratorStreamer
-import safetensors.torch
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import torch
+import safetensors.torch
 from fastapi.responses import StreamingResponse
 from tqdm import tqdm
-import logging
-import json
+import re
 
 load_dotenv()
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
 AWS_REGION = os.getenv("AWS_REGION")
@@ -35,16 +30,10 @@ s3_client = boto3.client(
 app = FastAPI()
 
 class DownloadModelRequest(BaseModel):
-    model_id: str
+    model_name: str
     pipeline_task: str
     input_text: str
 
-    @field_validator('model_id')
-    def validate_model_id(cls, value):
-        if not value:
-            raise ValueError("model_id cannot be empty")
-        return value
-
 class S3DirectStream:
     def __init__(self, bucket_name):
         self.s3_client = boto3.client(
@@ -57,204 +46,135 @@ class S3DirectStream:
 
     def stream_from_s3(self, key):
         try:
-            logger.info(f"Downloading {key} from S3...")
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            logger.info(f"Downloaded {key} from S3 successfully.")
             return response['Body']
         except self.s3_client.exceptions.NoSuchKey:
-            logger.error(f"File {key} not found in S3")
-            raise HTTPException(status_code=404, detail=f"File {key} not found in S3")
+            raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error al descargar de S3: {e}")
 
     def file_exists_in_s3(self, key):
         try:
             self.s3_client.head_object(Bucket=self.bucket_name, Key=key)
-            logger.info(f"File {key} exists in S3.")
             return True
-        except self.s3_client.exceptions.ClientError:
-            logger.info(f"File {key} does not exist in S3.")
-            return False
+        except self.s3_client.exceptions.ClientError as e:
+            if e.response['Error']['Code'] == '404':
+                return False
+            raise HTTPException(status_code=500, detail=f"Error al verificar archivo en S3: {e}")
 
     def load_model_from_stream(self, model_prefix):
         try:
-            logger.info(f"Loading model {model_prefix}...")
-            if self.file_exists_in_s3(f"{model_prefix}/config.json"):
-                logger.info(f"Model {model_prefix} found in S3. Loading...")
-                return self.load_model_from_existing_s3(model_prefix)
+            model_files = self.list_model_files(model_prefix)
+            if not model_files:
+                self.download_and_upload_to_s3(model_prefix)
+                return self.load_model_from_stream(model_prefix)
+
+            config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
+            config_data = config_stream.read().decode("utf-8")
+
+            model_path = f"{model_prefix}/model.safetensors"
+            if self.file_exists_in_s3(model_path):
+                model_stream = self.stream_from_s3(model_path)
+                model = AutoModelForCausalLM.from_config(config_data)
+                model.load_state_dict(safetensors.torch.load_stream(model_stream))
+            elif model_files:
+                model = AutoModelForCausalLM.from_config(config_data)
+                state_dict = {}
+                for file_name in model_files:
+                    file_stream = self.stream_from_s3(f"{model_prefix}/{file_name}")
+                    tmp = torch.load(file_stream, map_location="cpu")
+                    state_dict.update(tmp)
+                model.load_state_dict(state_dict)
+            else:
+                raise HTTPException(status_code=500, detail="Modelo no encontrado")
 
-            logger.info(f"Model {model_prefix} not found in S3. Downloading and uploading...")
-            self.download_and_upload_to_s3(model_prefix)
-            logger.info(f"Downloaded and uploaded {model_prefix}. Loading from S3...")
-            return self.load_model_from_stream(model_prefix)
+            return model
         except HTTPException as e:
-            logger.error(f"Error loading model: {e}")
-            return None
-
-    def load_model_from_existing_s3(self, model_prefix):
-        logger.info(f"Loading config for {model_prefix} from S3...")
-        config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
-        config_dict = json.load(config_stream)
-        config = AutoConfig.from_pretrained(model_prefix, **config_dict)
-        logger.info(f"Config loaded for {model_prefix}.")
+            raise
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo: {e}")
 
+    def list_model_files(self, model_prefix):
         try:
-            api = HfApi()
-            model_files = api.list_repo_files(model_prefix)
-            state_dict = {}
-            for file_info in model_files:
-                if file_info.rfilename.endswith(('.bin', '.safetensors')):
-                    file_url = api.download_file(model_prefix, file_info.rfilename)
-                    model_path = os.path.join(model_prefix, file_info.rfilename)
-                    logger.info(f"Downloading model file from {file_url} to {model_path} ...")
-                    with requests.get(file_url, stream=True) as response:
-                        if response.status_code == 200:
-                            try:
-                                model_stream = response.raw
-                                if model_path.endswith(".safetensors"):
-                                    shard_state = safetensors.torch.load_stream(model_stream)
-                                elif model_path.endswith(".bin"):
-                                    shard_state = torch.load(model_stream, map_location="cpu")
-                                state_dict.update(shard_state)
-                                logger.info(f"Downloaded and loaded model file {model_path}")
-                            except Exception as e:
-                                logger.exception(f"Error loading model file {model_path}: {e}")
-                                raise
-                        else:
-                            logger.error(f"Error downloading {file_url} with status code: {response.status_code}")
-                            raise HTTPException(status_code=500, detail=f"Error downloading model file from Hugging Face")
+          response = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=f"{model_prefix}/pytorch_model-")
+          model_files = []
+          if 'Contents' in response:
+            for obj in response['Contents']:
+              if re.match(r"pytorch_model-\d+-of-\d+", obj['Key'].split('/')[-1]):
+                model_files.append(obj['Key'].split('/')[-1])
+          return model_files
         except Exception as e:
-            logger.exception(f"Error loading model files for {model_prefix} : {e}")
-            raise
-
-
-        model = AutoModelForCausalLM.from_config(config)
-        model.load_state_dict(state_dict)
-        return model
+          return None
 
     def load_tokenizer_from_stream(self, model_prefix):
         try:
-            logger.info(f"Loading tokenizer for {model_prefix}...")
-            if self.file_exists_in_s3(f"{model_prefix}/tokenizer.json"):
-                logger.info(f"Tokenizer for {model_prefix} found in S3. Loading...")
-                return self.load_tokenizer_from_existing_s3(model_prefix, config)
-
-            logger.info(f"Tokenizer for {model_prefix} not found in S3. Downloading and uploading...")
-            self.download_and_upload_to_s3(model_prefix)
-            logger.info(f"Downloaded and uploaded tokenizer for {model_prefix}. Loading from S3...")
-            return self.load_tokenizer_from_stream(model_prefix)
+            tokenizer_path = f"{model_prefix}/tokenizer.json"
+            if self.file_exists_in_s3(tokenizer_path):
+                tokenizer_stream = self.stream_from_s3(tokenizer_path)
+                tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
+                return tokenizer
+            else:
+                self.download_and_upload_to_s3(model_prefix)
+                return self.load_tokenizer_from_stream(model_prefix)
         except HTTPException as e:
-            logger.error(f"Error loading tokenizer: {e}")
-            return None
-
-    def load_tokenizer_from_existing_s3(self, model_prefix, config):
-        logger.info(f"Loading tokenizer from S3 for {model_prefix}...")
-        tokenizer_stream = self.stream_from_s3(f"{model_prefix}/tokenizer.json")
-        tokenizer = AutoTokenizer.from_pretrained(None, config=config)
-        logger.info(f"Tokenizer loaded for {model_prefix}.")
-        return tokenizer
-
-    def download_and_upload_to_s3(self, model_prefix):
-        logger.info(f"Downloading and uploading model files for {model_prefix} to S3...")
-        try:
-            api = HfApi()
-            model_files = api.list_repo_files(model_prefix)
-
-            for file_info in model_files:
-                if file_info.rfilename.endswith(('.bin', '.safetensors', 'config.json', 'tokenizer.json')):
-                    file_url = api.download_file(model_prefix, file_info.rfilename)
-                    s3_key = f"{model_prefix}/{file_info.rfilename}"
-                    try:
-                        self.download_and_upload_to_s3_url(file_url, s3_key)
-                        logger.info(f"Downloaded and uploaded {s3_key}")
-                    except Exception as e:
-                        logger.exception(f"Error downloading/uploading {s3_key}: {e}")
-
-            logger.info(f"Finished downloading and uploading model files for {model_prefix}.")
-
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Error downloading model files from HuggingFace: {e}")
-            raise HTTPException(status_code=500, detail=f"Error downloading model files from Hugging Face") from e
+            raise
         except Exception as e:
-            logger.error(f"An unexpected error occurred: {e}")
-            raise HTTPException(status_code=500, detail=f"An unexpected error occurred during model download") from e
+            raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer: {e}")
 
 
-    def download_and_upload_to_s3_url(self, url, s3_key):
-        logger.info(f"Downloading from {url}...")
-        with requests.get(url, stream=True) as response:
-            if response.status_code == 200:
-                total_size_in_bytes = int(response.headers.get('content-length', 0))
-                block_size = 1024
-                progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
-                logger.info(f"Uploading to S3: {s3_key}...")
-                self.s3_client.upload_fileobj(response.raw, self.bucket_name, s3_key, Callback=lambda bytes_transferred: progress_bar.update(bytes_transferred))
-                progress_bar.close()
-                logger.info(f"Uploaded {s3_key} to S3 successfully.")
-            elif response.status_code == 404:
-                logger.error(f"File not found at {url}")
-                raise HTTPException(status_code=404, detail=f"Error downloading file from {url}. File not found.")
-            else:
-                logger.error(f"Error downloading from {url}: Status code {response.status_code}")
-                raise HTTPException(status_code=500, detail=f"Error downloading file from {url}")
-
-    def _get_latest_revision(self, model_prefix):
-        try:
-            api = HfApi()
-            model_info = api.model_info(model_prefix)
-            if hasattr(model_info, 'revision'):
-                revision = model_info.revision
-                if revision:
-                    return revision
-                else:
-                    logger.warning(f"No revision found for {model_prefix}, using 'main'")
-                    return "main"
-            else:
-                logger.warning(f"ModelInfo object for {model_prefix} does not have a 'revision' attribute, using 'main'")
-                return "main"
-        except Exception as e:
-            logger.error(f"Error getting latest revision for {model_prefix}: {e}")
-            return None
+    def download_and_upload_to_s3(self, model_prefix):
+        urls = {
+            "pytorch_model.bin": f"https://huggingface.co/{model_prefix}/resolve/main/pytorch_model.bin",
+            "model.safetensors": f"https://huggingface.co/{model_prefix}/resolve/main/model.safetensors",
+            "tokenizer.json": f"https://huggingface.co/{model_prefix}/resolve/main/tokenizer.json",
+            "config.json": f"https://huggingface.co/{model_prefix}/resolve/main/config.json"
+        }
+
+        for filename, url in urls.items():
+            try:
+                response = requests.get(url, stream=True)
+                response.raise_for_status()
+                self.s3_client.upload_fileobj(response.raw, self.bucket_name, f"{model_prefix}/{filename}")
+            except requests.exceptions.RequestException as e:
+                raise HTTPException(status_code=500, detail=f"Error al descargar {filename}: {e}")
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"Error al subir {filename} a S3: {e}")
 
 
 @app.post("/predict/")
 async def predict(model_request: DownloadModelRequest):
     try:
-        logger.info(f"Received request: Model={model_request.model_id}, Task={model_request.pipeline_task}, Input={model_request.input_text}")
-        model_id = model_request.model_id
-        task = model_request.pipeline_task
-        input_text = model_request.input_text
-
         streamer = S3DirectStream(S3_BUCKET_NAME)
-        logger.info("Loading model and tokenizer...")
-        model = streamer.load_model_from_stream(model_id)
-
-        if model is None:
-            logger.error(f"Failed to load model {model_id}")
-            raise HTTPException(status_code=500, detail=f"Failed to load model {model_id}")
+        model = streamer.load_model_from_stream(model_request.model_name)
+        tokenizer = streamer.load_tokenizer_from_stream(model_request.model_name)
 
-        tokenizer = streamer.load_tokenizer_from_stream(model_id)
-        logger.info("Model and tokenizer loaded.")
-
-        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering",  "summarization", "zero-shot-classification"]:
-            raise HTTPException(status_code=400, detail="Unsupported pipeline task")
+        task = model_request.pipeline_task
+        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", "text-to-speech", "text-to-image", "text-to-audio", "text-to-video"]:
+            raise HTTPException(status_code=400, detail="Pipeline task no soportado")
 
-        if task == "text-generation":
-            logger.info("Starting text generation...")
-            text_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-            inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
-            generation_kwargs = dict(inputs, streamer=text_streamer)
-            model.generate(**generation_kwargs)
-            logger.info("Text generation finished.")
-            return StreamingResponse(iter([tokenizer.decode(token) for token in text_streamer]), media_type="text/event-stream")
+        nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer)
+        input_text = model_request.input_text
+        outputs = nlp_pipeline(input_text)
+
+        if task in ["text-generation", "translation", "fill-mask", "sentiment-analysis", "question-answering"]:
+            return {"response": outputs}
+        elif task == "text-to-image":
+            s3_key = f"{model_request.model_name}/generated_image.png"
+            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png")
+        elif task == "text-to-audio":
+            s3_key = f"{model_request.model_name}/generated_audio.wav"
+            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="audio/wav")
+        elif task == "text-to-video":
+            s3_key = f"{model_request.model_name}/generated_video.mp4"
+            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4")
         else:
-            logger.info(f"Starting pipeline task: {task}...")
-            nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, device_map="auto", trust_remote_code=True)
-            outputs = nlp_pipeline(input_text)
-            logger.info(f"Pipeline task {task} finished.")
-            return {"result": outputs}
+            raise HTTPException(status_code=400, detail="Tipo de tarea desconocido")
 
+    except HTTPException as e:
+        raise
     except Exception as e:
-        logger.exception(f"Error processing request: {e}")
-        raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error inesperado: {str(e)}")
+
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)