Upload scorer_pred_local.py

scorer_pred_local.py

@@ -0,0 +1,94 @@
+# !/usr/bin/env python
+# -*- coding:utf-8 -*-
+
+import os
+import time
+import torch
+
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+
+if __name__ == "__main__":
+  import argparse
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--scorer-model-path', type=str, default="", help="file path", required=True)
+  parser.add_argument('--input-file-path', type=str, default="", help="file path", required=True)
+  parser.add_argument('--output-file-path', type=str, default="", help="file path", required=True)
+  parser.add_argument('--score-thres', type=float, default=3.0, help="score thres", required=False)
+  parser.add_argument('--text-key', type=str, default="text", help="file path", required=False)
+  parser.add_argument('--output-key', type=str, default="score", help="file path", required=False)
+  parser.add_argument('--do-score-filter', action='store_true', default=False, help='do score filter or not', dest='do_score_filter')
+  args = parser.parse_args()
+
+  # model_dir = '/share/project/ldwang/Aquila3/quality_scorer_base_from_qwen15_0_5b_labeled_by_deepspeek-v2'
+  model_dir = args.scorer_model_path
+  model = AutoModelForSequenceClassification.from_pretrained(
+    model_dir,
+    trust_remote_code=False,
+    ignore_mismatched_sizes=False,)
+  model.cuda()
+  model.eval()
+
+  tokenizer = AutoTokenizer.from_pretrained(
+    model_dir,
+    use_fast=True,
+    token=None,
+    trust_remote_code=False,)
+
+  max_length = 2048
+
+  import jsonlines
+  file_path = args.input_file_path
+  output_file_path = args.output_file_path
+  writer = jsonlines.open(output_file_path, mode='w')
+
+  dir_path = None
+  if os.path.isdir(file_path):
+    dir_path = os.listdir(file_path)
+  else:
+    dir_path = [file_path]
+
+  lines = 0
+  filtered = 0
+  start_time = time.time()
+
+  for file_path in dir_path:
+    input_file = os.path.join(args.input_file_path, file_path)
+    with jsonlines.open(input_file) as reader:
+      for line in reader:
+        lines += 1
+        if lines % 1000 == 0:
+          end_time = time.time()
+          elapsed_time = end_time - start_time
+          samples_per_second = lines / elapsed_time
+          print(f"Processed {lines} lines in {elapsed_time:.2f} seconds.", flush=True)
+          print(f"Samples per second: {samples_per_second:.2f}.", flush=True)
+
+        if args.text_key not in line:
+          filtered += 1
+          continue
+        sentecnce = line[args.text_key]
+        result = tokenizer(
+          [sentecnce],
+          padding=False,
+          max_length=max_length,
+          truncation=True,
+          return_tensors="pt",).to("cuda")
+        for key in result:
+          result[key] = torch.tensor(result[key])
+  
+        model_out = model(**result)
+        score = float(model_out.logits.tolist()[0][0])
+        if args.do_score_filter and score < args.score_thres:
+          filtered += 1
+          continue
+  
+        line[args.output_key] = score
+        writer.write(line)
+
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    samples_per_second = lines / elapsed_time
+    print(f"Processed {lines} lines in {elapsed_time:.2f} seconds, Filtered {filtered} samples.", flush=True)
+    print(f"Samples per second: {samples_per_second:.2f}.", flush=True)
+