handler without datasets

handler.py

@@ -1,7 +1,6 @@
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-from datasets import Dataset
 
 
 class EndpointHandler:
@@ -30,7 +29,7 @@ class EndpointHandler:
         """
         topics = data.pop("topics", data)
         texts = data.pop("texts", data)
-        batch_dict = {
+        batch = {
             'id': [],
             'text': [],
             'topic': []
@@ -38,16 +37,16 @@ class EndpointHandler:
 
         for topic in topics:
           for text in texts:        
-            batch_dict['id'].append(text['id'])
-            batch_dict['text'].append(text['text'])
-            batch_dict['topic'].append(topic)
-
-        batch = Dataset.from_dict(batch_dict)
+            batch['id'].append(text['id'])
+            batch['text'].append(text['text'])
+            batch['topic'].append(topic)
         
         tokenized_inputs = self.tokenize(batch)
 
         # run normal prediction
         output = self.model(**tokenized_inputs)
-        batch = batch.add_column('predictions', torch.argmax(output.logits, dim=-1).numpy(force=True))
-        batch = batch.map(lambda b: {'label': [self.model.config.id2label[p] for p in b['predictions']]}, batched=True, remove_columns=['text', 'predictions'])
-        return batch.to_dict()
+        predictions = torch.argmax(output.logits, dim=-1).numpy(force=True)
+        batch['label'] = [self.model.config.id2label[p] for p in predictions]
+        batch.pop('text')
+        
+        return batch