|
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 |
|
|
|
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.1 |
|
stop_sequences: list = [] |
|
|
|
@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 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) |
|
|
|
class TokenIteratorStreamer: |
|
def __init__(self): |
|
self.queue = asyncio.Queue() |
|
|
|
def put(self, value): |
|
self.queue.put_nowait(value) |
|
|
|
def end(self): |
|
self.queue.put_nowait(None) |
|
|
|
async def __aiter__(self): |
|
return self |
|
|
|
async def __anext__(self): |
|
value = await self.queue.get() |
|
if value is None: |
|
raise StopAsyncIteration |
|
return value |
|
|
|
@app.post("/generate") |
|
async def generate(request: GenerateRequest): |
|
model_name = request.model_name |
|
input_text = request.input_text |
|
task_type = request.task_type |
|
|
|
generation_params = request.model_dump( |
|
exclude_none=True, |
|
exclude={'model_name', 'input_text', 'task_type', 'stream', 'chunk_delay'} |
|
) |
|
|
|
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}") |
|
|
|
text_pipeline = pipeline(task_type, model=model_name, token=HUGGINGFACE_HUB_TOKEN, device_map="auto") |
|
token_streamer = TokenIteratorStreamer() |
|
|
|
def generate_on_thread(pipeline, input_text, streamer, generation_params): |
|
try: |
|
pipeline(input_text, |
|
max_new_tokens=int(1e9), |
|
return_full_text=False, |
|
streamer=streamer, |
|
**generation_params) |
|
finally: |
|
streamer.end() |
|
|
|
thread = Thread(target=generate_on_thread, args=(text_pipeline, input_text, token_streamer, generation_params)) |
|
thread.start() |
|
|
|
async def event_stream() -> AsyncIterator[str]: |
|
chunk_size = 20 |
|
tokens_buffer = [] |
|
async for token in token_streamer: |
|
tokens_buffer.append(token) |
|
if len(tokens_buffer) >= chunk_size: |
|
yield f"data: {json.dumps({'tokens': tokens_buffer})}\n\n" |
|
tokens_buffer = [] |
|
await asyncio.sleep(request.chunk_delay) |
|
if tokens_buffer: |
|
yield f"data: {json.dumps({'tokens': tokens_buffer})}\n\n" |
|
yield "\n\n" |
|
|
|
return StreamingResponse(event_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__": |
|
uvicorn.run(app, host="0.0.0.0", port=7860) |
