Add electra model

tasks/text.py

@@ -4,12 +4,15 @@ from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import random
 
+# Load model using Keras
+from tensorflow import keras
+
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "electra fine tune"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -22,6 +25,15 @@ async def evaluate_text(request: TextEvaluationRequest):
     - Makes random predictions from the label space (0-7)
     - Used as a baseline for comparison
     """
+    # Download from Google Drive
+    import gdown
+
+    url = 'https://drive.google.com/uc?id=1-HWE2G6ANbd7mqILdB9DPrvF3DrKI1e4'
+    output = 'checkpoint_epoch_5.weights.h5'
+    gdown.download(url, output, quiet=False)
+    
+    model = tf.keras.models.load_model('checkpoint_epoch_5.weights.h5')
+    
     # Get space info
     username, space_url = get_space_info()
 
@@ -57,13 +69,18 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
     
     # Make random predictions (placeholder for actual model inference)
+    # true_labels = test_dataset["label"]
+    # predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+
+    # 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
     #--------------------------------------------------------------------------------------------   
-
     
     # Stop tracking emissions
     emissions_data = tracker.stop_task()