import os from fastapi import FastAPI, HTTPException from fastapi.responses import StreamingResponse from pydantic import BaseModel, field_validator from transformers import pipeline, AutoConfig, AutoTokenizer from transformers.utils import logging from google.cloud import storage from google.auth.exceptions import DefaultCredentialsError import uvicorn import asyncio import json from huggingface_hub import login from dotenv import load_dotenv import huggingface_hub from threading import Thread from typing import AsyncIterator, List, Dict from transformers import StoppingCriteria, StoppingCriteriaList import torch load_dotenv() GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME") GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON") HUGGINGFACE_HUB_TOKEN = os.getenv("HF_API_TOKEN") if HUGGINGFACE_HUB_TOKEN: login(token=HUGGINGFACE_HUB_TOKEN) os.system("git config --global credential.helper store") if HUGGINGFACE_HUB_TOKEN: huggingface_hub.login(token=HUGGINGFACE_HUB_TOKEN, add_to_git_credential=True) logging.set_verbosity_info() logger = logging.get_logger(__name__) try: credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON) client = storage.Client.from_service_account_info(credentials_info) bucket = client.get_bucket(GCS_BUCKET_NAME) logger.info(f"Connection to Google Cloud Storage successful. Bucket: {GCS_BUCKET_NAME}") except (DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e: logger.error(f"Error loading credentials or bucket: {e}") raise RuntimeError(f"Error loading credentials or bucket: {e}") app = FastAPI() class GenerateRequest(BaseModel): model_name: str input_text: str task_type: str temperature: float = 1.0 stream: bool = True top_p: float = 1.0 top_k: int = 50 repetition_penalty: float = 1.0 num_return_sequences: int = 1 do_sample: bool = False chunk_delay: float = 0.0 max_new_tokens: int = 10 stopping_strings: List[str] = None @field_validator("model_name") def model_name_cannot_be_empty(cls, v): if not v: raise ValueError("model_name cannot be empty.") return v @field_validator("task_type") def task_type_must_be_valid(cls, v): valid_types = ["text-generation"] if v not in valid_types: raise ValueError(f"task_type must be one of: {valid_types}") return v class StopOnKeywords(StoppingCriteria): def __init__(self, stop_words_ids: List[List[int]], tokenizer, encounters: int = 1): super().__init__() self.stop_words_ids = stop_words_ids self.tokenizer = tokenizer self.encounters = encounters self.current_encounters = 0 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: for stop_ids in self.stop_words_ids: if torch.all(input_ids[0][-len(stop_ids):] == torch.tensor(stop_ids).to(input_ids.device)): self.current_encounters += 1 if self.current_encounters >= self.encounters: return True return False class GCSModelLoader: def __init__(self, bucket): self.bucket = bucket def _get_gcs_uri(self, model_name): return f"{model_name}" def _blob_exists(self, blob_path): blob = self.bucket.blob(blob_path) return blob.exists() def _create_model_folder(self, model_name): gcs_model_folder = self._get_gcs_uri(model_name) if not self._blob_exists(f"{gcs_model_folder}/.touch"): blob = self.bucket.blob(f"{gcs_model_folder}/.touch") blob.upload_from_string("") logger.info(f"Created folder '{gcs_model_folder}' in GCS.") def check_model_exists_locally(self, model_name): gcs_model_path = self._get_gcs_uri(model_name) blobs = self.bucket.list_blobs(prefix=gcs_model_path) return any(blobs) def download_model_from_huggingface(self, model_name): logger.info(f"Downloading model '{model_name}' from Hugging Face.") try: tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN) config = AutoConfig.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN) gcs_model_folder = self._get_gcs_uri(model_name) self._create_model_folder(model_name) tokenizer.save_pretrained(gcs_model_folder) config.save_pretrained(gcs_model_folder) for filename in os.listdir(config.name_or_path): if filename.endswith((".bin", ".safetensors")): blob = self.bucket.blob(f"{gcs_model_folder}/{filename}") blob.upload_from_filename(os.path.join(config.name_or_path, filename)) logger.info(f"Model '{model_name}' downloaded and saved to GCS.") return True except Exception as e: logger.error(f"Error downloading model from Hugging Face: {e}") return False model_loader = GCSModelLoader(bucket) @app.post("/generate") async def generate(request: GenerateRequest): model_name = request.model_name input_text = request.input_text task_type = request.task_type requested_max_new_tokens = request.max_new_tokens generation_params = request.model_dump( exclude_none=True, exclude={'model_name', 'input_text', 'task_type', 'stream', 'chunk_delay', 'max_new_tokens', 'stopping_strings'} ) user_defined_stopping_strings = request.stopping_strings try: if not model_loader.check_model_exists_locally(model_name): if not model_loader.download_model_from_huggingface(model_name): raise HTTPException(status_code=500, detail=f"Failed to load model: {model_name}") tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN) config = AutoConfig.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN) stopping_criteria_list = StoppingCriteriaList() if user_defined_stopping_strings: stop_words_ids = [tokenizer.encode(stop_string, add_special_tokens=False) for stop_string in user_defined_stopping_strings] stopping_criteria_list.append(StopOnKeywords(stop_words_ids, tokenizer)) # Pass tokenizer if config.eos_token_id is not None: eos_token_ids = [config.eos_token_id] if isinstance(config.eos_token_id, int): eos_token_ids = [[config.eos_token_id]] elif isinstance(config.eos_token_id, list): eos_token_ids = [[id] for id in config.eos_token_id] stop_words_ids_eos = [tokenizer.encode(tokenizer.decode(eos_id), add_special_tokens=False) for eos_id in eos_token_ids] stopping_criteria_list.append(StopOnKeywords(stop_words_ids_eos, tokenizer)) # Pass tokenizer elif tokenizer.eos_token is not None: stop_words_ids_eos = [tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)] stopping_criteria_list.append(StopOnKeywords(stop_words_ids_eos, tokenizer)) # Pass tokenizer async def generate_responses() -> AsyncIterator[Dict[str, List[Dict[str, str]]]]: nonlocal input_text all_generated_text = "" stop_reason = None while True: text_pipeline = pipeline( task_type, model=model_name, tokenizer=tokenizer, token=HUGGINGFACE_HUB_TOKEN, stopping_criteria=stopping_criteria_list, **generation_params, max_new_tokens=requested_max_new_tokens ) def generate_on_thread(pipeline, current_input_text, output_queue): result = pipeline(current_input_text) output_queue.put_nowait(result) output_queue = asyncio.Queue() thread = Thread(target=generate_on_thread, args=(text_pipeline, input_text, output_queue)) thread.start() result = await output_queue.get() thread.join() newly_generated_text = result[0]['generated_text'] # Decode tokens to check for stopping strings for criteria in stopping_criteria_list: if isinstance(criteria, StopOnKeywords): for stop_ids in criteria.stop_words_ids: decoded_stop_string = tokenizer.decode(stop_ids) if decoded_stop_string in newly_generated_text: stop_reason = f"stopping_string: {decoded_stop_string}" break if stop_reason: break if stop_reason: break all_generated_text += newly_generated_text yield {"response": [{'generated_text': newly_generated_text}]} if config.eos_token_id is not None: eos_tokens = [config.eos_token_id] if isinstance(config.eos_token_id, int): eos_tokens = [config.eos_token_id] elif isinstance(config.eos_token_id, list): eos_tokens = config.eos_token_id for eos_token in eos_tokens: if tokenizer.decode([eos_token]) in newly_generated_text: stop_reason = "eos_token" break if stop_reason: break elif tokenizer.eos_token is not None and tokenizer.eos_token in newly_generated_text: stop_reason = "eos_token" break input_text = all_generated_text async def text_stream(): async for data in generate_responses(): yield f"data: {json.dumps(data)}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(text_stream(), media_type="text/event-stream") except HTTPException as e: raise e except Exception as e: logger.error(f"Internal server error: {e}") raise HTTPException(status_code=500, detail=f"Internal server error: {e}") if __name__ == "__main__": import torch uvicorn.run(app, host="0.0.0.0", port=7860)