app.py

from huggingface_hub import HfApi
from fastapi import FastAPI, HTTPException
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
import torch
from fastapi.responses import StreamingResponse

load_dotenv()

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_name: str
    pipeline_task: str
    input_text: str
    revision: str = "main"

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:
            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
            return response['Body']
        except self.s3_client.exceptions.NoSuchKey:
            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)
            return True
        except self.s3_client.exceptions.ClientError:
            return False

    def load_model_from_stream(self, model_prefix, revision):
        try:
            if self.file_exists_in_s3(f"{model_prefix}/config.json") and \
               (self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin") or self.file_exists_in_s3(f"{model_prefix}/model.safetensors")):
                return self.load_model_from_existing_s3(model_prefix)

            self.download_and_upload_to_s3(model_prefix, revision)
            return self.load_model_from_stream(model_prefix, revision)
        except HTTPException as e:
            return None

    def load_model_from_existing_s3(self, model_prefix):
        config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
        config = AutoConfig.from_pretrained(config_stream) # Directly from stream

        if self.file_exists_in_s3(f"{model_prefix}/model.safetensors"):
            model_stream = self.stream_from_s3(f"{model_prefix}/model.safetensors")
            model = AutoModelForCausalLM.from_config(config)
            model.load_state_dict(safetensors.torch.load_stream(model_stream))
        elif self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin"):
            model_stream = self.stream_from_s3(f"{model_prefix}/pytorch_model.bin")
            model = AutoModelForCausalLM.from_config(config)
            state_dict = torch.load(model_stream, map_location="cpu")  # Load directly
            model.load_state_dict(state_dict)
        else:
           raise EnvironmentError(f"No model file found for {model_prefix} in S3")
        return model



    def load_tokenizer_from_stream(self, model_prefix):
        try:
            if self.file_exists_in_s3(f"{model_prefix}/tokenizer.json"):
                return self.load_tokenizer_from_existing_s3(model_prefix)
            self.download_and_upload_to_s3(model_prefix)
            return self.load_tokenizer_from_stream(model_prefix)
        except HTTPException as e:
            return None

    def load_tokenizer_from_existing_s3(self, model_prefix):
        tokenizer_stream = self.stream_from_s3(f"{model_prefix}/tokenizer.json")
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream) # Directly from stream
        return tokenizer


    def download_and_upload_to_s3(self, model_prefix, revision="main"):
        model_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/pytorch_model.bin"
        safetensors_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/model.safetensors"
        tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/tokenizer.json"
        config_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/config.json"

        self.download_and_upload_to_s3_url(model_url, f"{model_prefix}/pytorch_model.bin")
        self.download_and_upload_to_s3_url(safetensors_url, f"{model_prefix}/model.safetensors")
        self.download_and_upload_to_s3_url(tokenizer_url, f"{model_prefix}/tokenizer.json")
        self.download_and_upload_to_s3_url(config_url, f"{model_prefix}/config.json")


    def download_and_upload_to_s3_url(self, url, s3_key):
        response = requests.get(url, stream=True)
        if response.status_code == 200:
            self.s3_client.upload_fileobj(response.raw, self.bucket_name, s3_key) # Direct upload
        elif response.status_code == 404:
            raise HTTPException(status_code=404, detail=f"Error downloading file from {url}. File not found.")

        else:
            raise HTTPException(status_code=500, detail=f"Error downloading file from {url}")



@app.post("/predict/")
async def predict(model_request: DownloadModelRequest):
    try:
        model_name = model_request.model_name
        revision = model_request.revision

        streamer = S3DirectStream(S3_BUCKET_NAME)
        model = streamer.load_model_from_stream(model_name, revision)
        tokenizer = streamer.load_tokenizer_from_stream(model_name)

        task = model_request.pipeline_task
        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":
            text_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
            inputs = tokenizer(model_request.input_text, return_tensors="pt").to(model.device)
            generation_kwargs = dict(inputs, streamer=text_streamer)
            model.generate(**generation_kwargs)
            return StreamingResponse(iter([tokenizer.decode(token) for token in text_streamer]), media_type="text/event-stream")

        else:
            nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, device_map="auto", trust_remote_code=True)
            outputs = nlp_pipeline(model_request.input_text)
            return {"result": outputs}


    except Exception as e:
        print(f"Complete Error: {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)