Update app.py

app.py

@@ -24,6 +24,7 @@ import logging
 
 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")
 
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -88,15 +89,36 @@ class GCSModelLoader:
             return json.loads(blob.download_as_bytes())
         except Exception as e:
             raise Exception(f"Error downloading config: {e}")
+    
+    async def _upload_to_gcs(self, gcs_uri, model_files):
+        try:
+            for file_name, file_content in model_files.items():
+              blob_path = gcs_uri.replace(f'gs://{self.bucket.name}/', '') + '/' + file_name
+              blob = self.bucket.blob(blob_path)
+              blob.upload_from_string(file_content)
+        except Exception as e:
+            raise Exception(f"Error uploading to GCS: {e}")
 
     async def load_config(self, model_name):
-      gcs_uri = self._get_gcs_uri(model_name)
-      try:
-        config_data = await self._download_config_from_gcs(gcs_uri)
-        config = AutoConfig.from_dict(config_data)
-        return config
-      except Exception as e:
-          raise HTTPException(status_code=500, detail=f"Error loading config: {e}")
+        gcs_uri = self._get_gcs_uri(model_name)
+        try:
+            config_data = await self._download_config_from_gcs(gcs_uri)
+            config = AutoConfig.from_pretrained(gcs_uri, config = config_data)
+            return config
+        except Exception as e:
+            try:
+              config = AutoConfig.from_pretrained(model_name, use_auth_token=HUGGINGFACE_HUB_TOKEN)
+              tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HUGGINGFACE_HUB_TOKEN)
+              model_files = {}
+              config_file = f"{model_name}/config.json"
+              tokenizer_file = f"{model_name}/tokenizer.json"
+              model_files[config_file] = json.dumps(config.to_dict())
+              model_files[tokenizer_file] = json.dumps(tokenizer.to_dict())
+              await self._upload_to_gcs(gcs_uri, model_files)
+              return config
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"Error loading config: {e}")
+
 
 model_loader = GCSModelLoader(bucket, client)
 
@@ -130,7 +152,7 @@ async def stream_from_gcs(model_name, input_text, generation_config, stop_sequen
         yield str(token)
     del model
 
-class CustomTextIteratorStreamer:
+class CustomTextIteratorStreamer(TextIteratorStreamer):
     def __init__(self, gcs_uri, chunk_delay):
         self.chunk_delay = chunk_delay
         self.queue = asyncio.Queue()