Update app.py

app.py

@@ -7,7 +7,6 @@ from transformers import (
     GenerationConfig,
     AutoTokenizer,
     AutoModelForCausalLM,
-    TextIteratorStreamer
 )
 from google.cloud import storage
 from google.auth.exceptions import DefaultCredentialsError
@@ -80,7 +79,9 @@ class GCSModelLoader:
 
     def _download_content(self, blob_path):
         blob = self.bucket.blob(blob_path)
-        return blob.download_as_bytes() if blob.exists() else None
+        if blob.exists():
+            return blob.download_as_bytes()
+        return None
 
     def _upload_content(self, content, blob_path):
         blob = self.bucket.blob(blob_path)
@@ -99,8 +100,7 @@ class GCSModelLoader:
             try:
                 config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
-                config_json = json.dumps(config.to_dict())
-                self._upload_content(config_json.encode('utf-8'), f"{gcs_model_folder}/config.json")
+                self._upload_content(json.dumps(config.to_dict()).encode('utf-8'), f"{gcs_model_folder}/config.json")
                 return config
             except Exception as e:
                 logger.error(f"Error loading config from Hugging Face and saving to GCS: {e}")
@@ -121,9 +121,10 @@ 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 file in tokenizer.save_pretrained(None):
-                    with open(file, 'rb') as f:
-                        self._upload_content(f.read(), f"{gcs_model_folder}/{os.path.basename(file)}")
+                os.makedirs(gcs_model_folder, exist_ok=True)
+                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}")
                 return tokenizer
             except Exception as e:
                 logger.error(f"Error loading tokenizer from Hugging Face and saving to GCS: {e}")
@@ -144,6 +145,7 @@ 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)
+                os.makedirs(gcs_model_folder, exist_ok=True)
                 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}")
@@ -154,17 +156,20 @@ class GCSModelLoader:
 
 model_loader = GCSModelLoader(bucket)
 
-async def generate_stream(model, tokenizer, input_text, generation_config, chunk_delay):
+async def generate_stream(model, tokenizer, input_text, generation_config):
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = dict(inputs, generation_config=generation_config, streamer=streamer)
-    asyncio.create_task(model.generate(**generation_kwargs))
-
-    async def event_stream():
-        for token in streamer:
+    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(chunk_delay)
-    return event_stream()
+            await asyncio.sleep(0.001) # Adjust delay as needed
+    return token_stream()
 
 def generate_non_stream(model, tokenizer, input_text, generation_config):
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
@@ -178,19 +183,15 @@ async def generate(request: GenerateRequest):
     task_type = request.task_type
     stream = request.stream
 
-    generation_params = {}
-    for key, value in request.model_dump(exclude_none=True, exclude={'model_name', 'input_text', 'task_type', 'stream', 'chunk_delay'}).items():
-        generation_params[key] = value
+    generation_params = request.model_dump(
+        exclude_none=True,
+        exclude={'model_name', 'input_text', 'task_type', 'stream', 'chunk_delay'}
+    )
 
     try:
         gcs_model_folder_uri = model_loader._get_gcs_uri(model_name)
-        if not bucket.blob(f"{gcs_model_folder_uri}/config.json").exists():
-            logger.info(f"Model '{model_name}' not found in GCS, creating placeholder and downloading.")
-            bucket.blob(f"{gcs_model_folder_uri}/.placeholder").upload_from_string("")
-            for folder in gcs_model_folder_uri.split('/'):
-                prefix = "/".join(gcs_model_folder_uri.split('/')[:gcs_model_folder_uri.split('/').index(folder)+1])
-                if not any(bucket.list_blobs(prefix=f"{prefix}/config.json")):
-                    bucket.blob(f"{prefix}/.placeholder").upload_from_string("")
+        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.")
 
         config = model_loader.load_config(model_name)
         if not config:
@@ -222,12 +223,12 @@ async def generate(request: GenerateRequest):
 
         if task_type == "text-to-text":
             if stream:
-                return StreamingResponse(generate_stream(model, tokenizer, input_text, generation_config, request.chunk_delay), media_type="text/event-stream")
+                return StreamingResponse(generate_stream(model, tokenizer, input_text, generation_config), media_type="text/event-stream")
             else:
                 text_result = generate_non_stream(model, tokenizer, input_text, generation_config)
                 return {"text": text_result}
         else:
-             raise HTTPException(status_code=400, detail=f"Task type not supported: {task_type}")
+            raise HTTPException(status_code=400, detail=f"Task type not supported: {task_type}")
 
     except HTTPException as e:
         raise e