app.py

from huggingface_hub import HfApi
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, field_validator
import requests
import boto3
from dotenv import load_dotenv
import os
import uvicorn
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoConfig, TextIteratorStreamer
import safetensors.torch
import torch
from fastapi.responses import StreamingResponse
from tqdm import tqdm
import logging
import json

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")
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

s3_client = boto3.client(
    's3',
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
    region_name=AWS_REGION
)

app = FastAPI()

class DownloadModelRequest(BaseModel):
    model_id: 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(
            's3',
            aws_access_key_id=AWS_ACCESS_KEY_ID,
            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
            region_name=AWS_REGION
        )
        self.bucket_name = bucket_name

    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")

    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

    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") and \
               any(self.file_exists_in_s3(f"{model_prefix}/{file}") for file in self._get_model_files(model_prefix)):
                logger.info(f"Model {model_prefix} found in S3. Loading...")
                return self.load_model_from_existing_s3(model_prefix)

            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)
        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}.")

        model_files = self._get_model_files(model_prefix)
        if not model_files:
            logger.error(f"No model files found for {model_prefix} in S3")
            raise EnvironmentError(f"No model files found for {model_prefix} in S3")

        state_dict = {}
        for model_file in model_files:
            model_path = os.path.join(model_prefix, model_file)
            logger.info(f"Loading model file: {model_path}")
            model_stream = self.stream_from_s3(model_path)
            try:
                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")
                else:
                    logger.error(f"Unsupported model file type: {model_path}")
                    raise ValueError(f"Unsupported model file type: {model_path}")

                state_dict.update(shard_state)
            except Exception as e:
                logger.exception(f"Error loading model file {model_path}: {e}")
                raise

        model = AutoModelForCausalLM.from_config(config)
        model.load_state_dict(state_dict)
        return model

    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)
        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, revision="main"):
        logger.info(f"Downloading and uploading model files for {model_prefix} to S3...")
        config_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/config.json"
        self.download_and_upload_to_s3_url(config_url, f"{model_prefix}/config.json")

        model_files = self._get_model_files(model_prefix, revision)
        for model_file in model_files:
            url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/{model_file}"
            s3_key = f"{model_prefix}/{model_file}"
            self.download_and_upload_to_s3_url(url, s3_key)
            logger.info(f"Downloaded and uploaded {s3_key}")

        tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/tokenizer.json"
        self.download_and_upload_to_s3_url(tokenizer_url, f"{model_prefix}/tokenizer.json")

        logger.info(f"Finished downloading and uploading model files for {model_prefix}.")


    def _get_model_files(self, model_prefix, revision="main"):
        index_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/"
        try:
            index_response = requests.get(index_url)
            index_response.raise_for_status()
            logger.info(f"Hugging Face API Response: Status Code = {index_response.status_code}, Headers = {index_response.headers}")
            index_content = index_response.text
            logger.info(f"Index content: {index_content}")
            model_files = [f for f in index_content.split('\n') if f.endswith(('.bin', '.safetensors'))]
            return model_files
        except requests.exceptions.RequestException as e:
            logger.error(f"Error retrieving model index: {e}")
            raise HTTPException(status_code=500, detail=f"Error retrieving model files from Hugging Face") from e
        except (IndexError, ValueError) as e:
            logger.error(f"Error parsing model file names from Hugging Face: {e}")
            raise HTTPException(status_code=500, detail=f"Error retrieving model files from Hugging Face") from 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


@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}")

        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")

        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")
        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}

    except Exception as e:
        logger.exception(f"Error processing request: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)