Added some efficiency

tasks/text.py

@@ -18,6 +18,31 @@ router = APIRouter()
 DESCRIPTION = "Electra_Base"
 ROUTE = "/text"
 
+class CustomTFDataset(tf.data.Dataset):
+    def __init__(self, texts, labels, tokenizer, max_length=128):
+        self.texts = texts
+        self.labels = labels
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __iter__(self):
+        for text, label in zip(self.texts, self.labels):
+            encoding = self.tokenizer(
+                text,
+                truncation=True,
+                padding='max_length',
+                max_length=self.max_length,
+                return_tensors='tf'
+            )
+            yield {
+                'input_ids': encoding['input_ids'][0],
+                'attention_mask': encoding['attention_mask'][0],
+                'label': tf.constant(label, dtype=tf.int32)
+            }
+
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):