Update tasks/text.py

tasks/text.py

@@ -7,20 +7,31 @@ import random
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+#modified: additional lib
+import tensorflow as tf
+from huggingface_hub import hf_hub_download
+from transformers import TFElectraForSequenceClassification, ElectraTokenizer, ElectraConfig
+#
+
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "Finetuned ELECTRA"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
+
+#modified: retrieve model
+model_repo = "jennasparks/electra_tf"
+config = ElectraConfig.from_pretrained(model_repo)
+model = TFElectraForSequenceClassification.from_pretrained(model_repo)
+tokenizer = ElectraTokenizer.from_pretrained("google/electra-base-discriminator")
+#
+
 async def evaluate_text(request: TextEvaluationRequest):
     """
     Evaluate text classification for climate disinformation detection.
-    
-    Current Model: Random Baseline
-    - Makes random predictions from the label space (0-7)
-    - Used as a baseline for comparison
+    Current Model: Finetuned ELECTRA
     """
     # Get space info
     username, space_url = get_space_info()
@@ -56,9 +67,16 @@ 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 = []
+    
+    for i in range(len(test_dataset["quote"])):
+    encoded_input = tokenizer(test_dataset["quote"][i], truncation=True, padding=True, return_tensors="tf")
+    outputs = model(encoded_input["input_ids"], attention_mask=encoded_input["attention_mask"], training=False)
+    predictions.append(tf.argmax(outputs.logits, axis=1))
+    
+    # Get true labels
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE