Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import tarfile
 import wandb
 
@@ -10,6 +11,7 @@ from transformers import ViTFeatureExtractor
 PRETRAIN_CHECKPOINT = "google/vit-base-patch16-224-in21k"
 feature_extractor = ViTFeatureExtractor.from_pretrained(PRETRAIN_CHECKPOINT)
 
+WB_KEY = os.environ['WB_KEY']
 MODEL = None 
 
 RESOLTUION = 224
@@ -43,11 +45,11 @@ def preprocess_input(image):
         "pixel_values": tf.expand_dims(image, 0)
     }
 
-def get_predictions(wb_token, image):
+def get_predictions(image):
     global MODEL
     
     if MODEL is None:
-        wandb.login(key=wb_token)
+        wandb.login(key=WB_KEY)
         wandb.init(project="tfx-vit-pipeline", id="gvtyqdgn", resume=True)
         path = wandb.use_artifact('tfx-vit-pipeline/final_model:1688113391', type='model').download()
 
@@ -66,8 +68,6 @@ def get_predictions(wb_token, image):
 with gr.Blocks() as demo:
     gr.Markdown("## Simple demo for a Image Classification of the Beans Dataset with HF ViT model")
 
-    wb_token_if = gr.Textbox(interactive=True, label="Your Weight & Biases API Key")
-
     with gr.Row():
         image_if = gr.Image()
         label_if = gr.Label(num_top_classes=3)
@@ -76,7 +76,7 @@ with gr.Blocks() as demo:
 
     classify_if.click(
         get_predictions,
-        [wb_token_if, image_if],
+        image_if,
         label_if
     )