Update infer.py

infer.py

@@ -147,25 +147,23 @@ def get_model_and_tokenizer(model_name: str, optimization_level: str, progress):
     )
 
 
-# def collate_fn(examples, tokenizer=None, padding=None, column_name="text"):
-#     try:
-#         keys = examples[0].keys()
-#     except KeyError:
-#         print(examples)
-#     else:
-#         batch = {k: [] for k in examples[0].keys()}
-
-#     tokenized = tokenizer(
-#         [x[column_name] for x in examples],
-#         truncation=True,
-#         padding=padding,
-#         max_length=512,
-#         return_tensors="pt"
-#         )
+def collate_fn(examples, column_name, tokenizer):
+    feature_cols = ["input_ids", "attention_mask"]
+    features = [{k: x[k] for k in feature_cols} for x in examples]
+
+    print(features)
+
+    tokenized = tokenizer.pad(
+        features,
+        padding=True,
+        max_length=512,
+        return_tensors="pt",
+        pad_to_multiple_of=16,
+        )
 
-#     tokenized[column_name] = [x[column_name] for x in examples]
+    tokenized[column_name] = [x[column_name] for x in examples]
 
-#     return tokenized
+    return tokenized
 
 
 @torch.inference_mode()
@@ -247,8 +245,8 @@ def batch_embed(
 
     start_time = time.time()
 
-    collator = DataCollatorWithPadding(
-        tokenizer, padding=True, max_length=512, pad_to_multiple_of=16
+    collator = partial(
+        collate_fn, column_name=column_name, tokenizer=tokenizer
     )
 
     dl = DataLoader(