Update app.py

app.py

@@ -15,6 +15,10 @@ import asyncio
 import json
 import logging
 from huggingface_hub import login
+from dotenv import load_dotenv
+import huggingface_hub
+
+load_dotenv()
 
 GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
 GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
@@ -23,6 +27,9 @@ 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")
+huggingface_hub.login(token=HUGGINGFACE_HUB_TOKEN, add_to_git_credential=True)
+
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
@@ -71,7 +78,7 @@ class GCSModelLoader:
         self.bucket = bucket
 
     def _get_gcs_uri(self, model_name):
-         return f"{model_name}"
+        return f"{model_name}"
 
     def _blob_exists(self, blob_path):
         blob = self.bucket.blob(blob_path)
@@ -87,20 +94,27 @@ class GCSModelLoader:
         blob = self.bucket.blob(blob_path)
         blob.upload_from_string(content)
 
+    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 load_config(self, model_name):
-      
         gcs_config_path = f"{self._get_gcs_uri(model_name)}/config.json"
         config_content = self._download_content(gcs_config_path)
         if config_content:
             try:
                 return AutoConfig.from_pretrained(pretrained_model_name_or_path=None, trust_remote_code=True, config_dict=json.loads(config_content), token=HUGGINGFACE_HUB_TOKEN)
             except Exception as e:
-                 logger.error(f"Error loading config from GCS: {e}")
-                 return None
+                logger.error(f"Error loading config from GCS: {e}")
+                return None
         else:
             try:
                 config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
+                self._create_model_folder(model_name)
                 self._upload_content(json.dumps(config.to_dict()).encode('utf-8'), f"{gcs_model_folder}/config.json")
                 return config
             except Exception as e:
@@ -114,7 +128,7 @@ class GCSModelLoader:
 
         if gcs_files_exist:
             try:
-                return AutoTokenizer.from_pretrained(gcs_tokenizer_path, trust_remote_code=True,token=HUGGINGFACE_HUB_TOKEN)
+                return AutoTokenizer.from_pretrained(gcs_tokenizer_path, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
             except Exception as e:
                 logger.error(f"Error loading tokenizer from GCS: {e}")
                 return None
@@ -122,9 +136,8 @@ class GCSModelLoader:
             try:
                 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
-                for filename in os.listdir(tokenizer.save_pretrained(None)):
-                   with open(filename, 'rb') as f:
-                       self._upload_content(f.read(), f"{gcs_model_folder}/{filename}")
+                self._create_model_folder(model_name)
+                tokenizer.save_pretrained(gcs_model_folder)
                 return tokenizer
             except Exception as e:
                 logger.error(f"Error loading tokenizer from Hugging Face and saving to GCS: {e}")
@@ -145,9 +158,8 @@ class GCSModelLoader:
             try:
                 model = AutoModelForCausalLM.from_pretrained(model_name, config=config, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
-                for filename in os.listdir(model.save_pretrained(None)):
-                   with open(filename, 'rb') as f:
-                       self._upload_content(f.read(), f"{gcs_model_folder}/{filename}")
+                self._create_model_folder(model_name)
+                model.save_pretrained(gcs_model_folder)
                 return model
             except Exception as e:
                 logger.error(f"Error loading model from Hugging Face and saving to GCS: {e}")
@@ -157,19 +169,19 @@ model_loader = GCSModelLoader(bucket)
 
 async def generate_stream(model, tokenizer, input_text, generation_config):
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
-    generation_stream = model.generate(
-        **inputs,
-        generation_config=generation_config,
-        stream=True,
-    )
+
     async def token_stream():
-        for output in generation_stream:
-           token_id = output[-1]
-           token = tokenizer.decode(token_id, skip_special_tokens=True)
-           yield {"token": token}
-           await asyncio.sleep(0.001)
-    return token_stream()
+        generation_stream = model.generate(
+            **inputs,
+            generation_config=generation_config,
+            stream=True,
+        )
+        async for output in generation_stream:
+            token_id = output[-1]
+            token = tokenizer.decode(token_id, skip_special_tokens=True)
+            yield {"token": token}
 
+    return token_stream()
 
 def generate_non_stream(model, tokenizer, input_text, generation_config):
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
@@ -191,7 +203,7 @@ async def generate(request: GenerateRequest):
     try:
         gcs_model_folder_uri = model_loader._get_gcs_uri(model_name)
         if not model_loader._blob_exists(f"{gcs_model_folder_uri}/config.json"):
-           logger.info(f"Model '{model_name}' not found in GCS, downloading from Hugging Face.")
+            logger.info(f"Model '{model_name}' not found in GCS, checking Hugging Face.")
 
         config = model_loader.load_config(model_name)
         if not config:
@@ -199,21 +211,17 @@ async def generate(request: GenerateRequest):
 
         tokenizer = model_loader.load_tokenizer(model_name)
         if not tokenizer:
-           raise HTTPException(status_code=400, detail="Tokenizer could not be loaded.")
+            raise HTTPException(status_code=400, detail="Tokenizer could not be loaded.")
 
         generation_config_kwargs = generation_params.copy()
-        if hasattr(tokenizer, 'pad_token_id') and tokenizer.pad_token_id is not None:
-            generation_config_kwargs['pad_token_id'] = tokenizer.pad_token_id
-        if hasattr(tokenizer, 'eos_token_id') and tokenizer.eos_token_id is not None:
-            generation_config_kwargs['eos_token_id'] = tokenizer.eos_token_id
-        if hasattr(tokenizer, 'sep_token_id') and tokenizer.sep_token_id is not None:
-            generation_config_kwargs['sep_token_id'] = tokenizer.sep_token_id
-        if hasattr(tokenizer, 'unk_token_id') and tokenizer.unk_token_id is not None:
-            generation_config_kwargs['unk_token_id'] = tokenizer.unk_token_id
-        
+        generation_config_kwargs['pad_token_id'] = tokenizer.pad_token_id
+        generation_config_kwargs['eos_token_id'] = tokenizer.eos_token_id
+        generation_config_kwargs['sep_token_id'] = tokenizer.sep_token_id
+        generation_config_kwargs['unk_token_id'] = tokenizer.unk_token_id
+
         model = model_loader.load_model(model_name, config)
         if not model:
-           raise HTTPException(status_code=400, detail="Model could not be loaded.")
+            raise HTTPException(status_code=400, detail="Model could not be loaded.")
 
         generation_config = GenerationConfig.from_pretrained(
             model_name,
@@ -223,7 +231,11 @@ async def generate(request: GenerateRequest):
 
         if task_type == "text-to-text":
             if stream:
-                return StreamingResponse(generate_stream(model, tokenizer, input_text, generation_config), media_type="text/event-stream")
+                async def event_stream():
+                    async for output in generate_stream(model, tokenizer, input_text, generation_config):
+                        yield f"data: {json.dumps(output)}\n\n"
+                        await asyncio.sleep(request.chunk_delay)
+                return StreamingResponse(event_stream(), media_type="text/event-stream")
             else:
                 text_result = generate_non_stream(model, tokenizer, input_text, generation_config)
                 return {"text": text_result}