loading model from hf

tasks/text.py

@@ -8,13 +8,11 @@ from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 import tensorflow as tf
-
-# Download from Google Drive
-import gdown
+from huggingface_hub import hf_hub_download
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "Electra"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -42,10 +40,11 @@ async def evaluate_text(request: TextEvaluationRequest):
         "7_fossil_fuels_needed": 7
     }
 
-    # Download our pre-trained model
-    url = 'https://drive.google.com/uc?id=1-HWE2G6ANbd7mqILdB9DPrvF3DrKI1e4'
-    output = 'checkpoint_epoch_5.weights.h5'
-    gdown.download(url, output, quiet=False)
+    # Download our pre-trained model from Hugging Face
+    model_path = hf_hub_download(repo_id="jennasparks/frugal-ai-text-electra-base", filename="checkpoint_epoch_5.weights.h5")
+
+    # Load the model
+    model = tf.keras.models.load_model(model_path)
 
     # Load and prepare the dataset
     dataset = load_dataset(request.dataset_name)
@@ -66,9 +65,11 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
     #--------------------------------------------------------------------------------------------   
     
-    # Make random predictions (placeholder for actual model inference)
+    # Make predictions
+    predictions = model.predict(test_dataset)
+
+    # Get true labels
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE