app.py

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)