Update tasks/text.py

tasks/text.py

@@ -57,8 +57,45 @@ 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))]
+
+    from transformers import DistilBertTokenizer
+    import numpy as np
+    import onnxruntime as ort
+    
+    # Load the ONNX model and tokenizer
+    MODEL_PATH = "/Users/hinabandukwala/Documents/frugalai/submission-template/models/distilbert_quantized_dynamic.onnx"
+    tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+    ort_session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
+
+    # Preprocess the text data
+    def preprocess(texts):
+        return tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            max_length=365,
+            return_tensors="np"
+        )
+
+    # Run inference
+    def predict(texts):
+        inputs = preprocess(texts)
+        ort_inputs = {
+            "input_ids": inputs["input_ids"].astype(np.int64),
+            "attention_mask": inputs["attention_mask"].astype(np.int64)
+        }
+        ort_outputs = ort_session.run(None, ort_inputs)
+        logits = ort_outputs[0]
+        predictions = np.argmax(logits, axis=1)
+        return predictions
+
+     # Replace the random predictions with actual model predictions
+    texts = test_dataset["text"]
+    predictions = predict(texts)
+
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE