File size: 3,161 Bytes
a70a7f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 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 = 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)
|