|
|
|
|
|
|
|
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) |
|
|
|
|
