added inference

tasks/text.py

@@ -7,11 +7,19 @@ import random
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+
+import tensorflow as tf
+from huggingface_hub import hf_hub_download
+from transformers import ElectraTokenizer
+
+
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "Electra with balanced dataset"
 ROUTE = "/text"
 
+
+
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
@@ -37,6 +45,12 @@ async def evaluate_text(request: TextEvaluationRequest):
         "7_fossil_fuels_needed": 7
     }
 
+    # Download our pre-trained model from Hugging Face
+    model_path = hf_hub_download(repo_id="julianaconsuegra/electra-base-climate-disinformation", filename="tf_model.h5")
+
+    # Load the model
+    model = tf.keras.models.load_model(model_path)
+
     # Load and prepare the dataset
     dataset = load_dataset(request.dataset_name)
 
@@ -55,10 +69,26 @@ async def evaluate_text(request: TextEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # 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.
     #--------------------------------------------------------------------------------------------   
+    # Load ELECTRA tokenizer
+    tokenizer = ElectraTokenizer.from_pretrained("google/electra-base-discriminator")
+    
+    # Tokenize test data with same parameters as training
+    inputs = tokenizer(
+        test_dataset["text"],
+        truncation=True,
+        padding="max_length",
+        return_tensors="tf"
+    )
+    
+    # Run model prediction
+    logits = model.predict({
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"]
+    })
+    predictions = tf.argmax(logits, axis=1).numpy()
     
-    # Make random predictions (placeholder for actual model inference)
+    # Get ground truth labels
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE