Spaces:
Sleeping
Sleeping
| 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 | |
| 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}...") | |
| revision = self._get_latest_revision(model_prefix) | |
| if revision is None: | |
| logger.error(f"Could not determine revision for {model_prefix}") | |
| raise ValueError(f"Could not determine revision for {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, revision)): | |
| 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, revision) | |
| 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) | |
| revision = model_info.revision | |
| if revision: | |
| return revision | |
| else: | |
| logger.warning(f"No default revision found for {model_prefix}, using 'main'") | |
| return "main" | |
| except Exception as e: | |
| logger.error(f"Error getting latest revision for {model_prefix}: {e}") | |
| return None | |
| 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) |