|
|
@@ -1,86 +0,0 @@
|
|
|
-import gzip
|
|
|
-import json
|
|
|
-
|
|
|
-import torch
|
|
|
-import typer
|
|
|
-from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu
|
|
|
-from tqdm import tqdm
|
|
|
-from transformers import LayoutLMv3ForTokenClassification
|
|
|
-
|
|
|
-from helpers import (
|
|
|
- DataCollator,
|
|
|
- check_duplicate,
|
|
|
- MAX_LEN,
|
|
|
- parse_logits,
|
|
|
- prepare_inputs,
|
|
|
-)
|
|
|
-
|
|
|
-app = typer.Typer()
|
|
|
-
|
|
|
-chen_cherry = SmoothingFunction()
|
|
|
-
|
|
|
-
|
|
|
-@app.command()
|
|
|
-def main(
|
|
|
- input_file: str = typer.Argument(..., help="input file"),
|
|
|
- model_path: str = typer.Argument(..., help="model path"),
|
|
|
- batch_size: int = typer.Option(16, help="batch size"),
|
|
|
-):
|
|
|
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
- model = (
|
|
|
- LayoutLMv3ForTokenClassification.from_pretrained(model_path, num_labels=MAX_LEN)
|
|
|
- .bfloat16()
|
|
|
- .to(device)
|
|
|
- .eval()
|
|
|
- )
|
|
|
- data_collator = DataCollator()
|
|
|
- if torch.cuda.is_available():
|
|
|
- torch.cuda.empty_cache()
|
|
|
-
|
|
|
- datasets = []
|
|
|
- with gzip.open(input_file, "rt") as f:
|
|
|
- for line in tqdm(f):
|
|
|
- datasets.append(json.loads(line))
|
|
|
- # make batch faster
|
|
|
- datasets.sort(key=lambda x: len(x["source_boxes"]), reverse=True)
|
|
|
-
|
|
|
- total = 0
|
|
|
- total_out_idx = 0.0
|
|
|
- total_out_token = 0.0
|
|
|
- for i in tqdm(range(0, len(datasets), batch_size)):
|
|
|
- batch = datasets[i : i + batch_size]
|
|
|
- model_inputs = data_collator(batch)
|
|
|
- model_inputs = prepare_inputs(model_inputs, model)
|
|
|
- # forward
|
|
|
- with torch.no_grad():
|
|
|
- model_outputs = model(**model_inputs)
|
|
|
- logits = model_outputs.logits.cpu()
|
|
|
- for data, logit in zip(batch, logits):
|
|
|
- target_index = data["target_index"][:MAX_LEN]
|
|
|
- pred_index = parse_logits(logit, len(target_index))
|
|
|
- assert len(pred_index) == len(target_index)
|
|
|
- assert not check_duplicate(pred_index)
|
|
|
- target_texts = data["target_texts"][:MAX_LEN]
|
|
|
- source_texts = data["source_texts"][:MAX_LEN]
|
|
|
- pred_texts = []
|
|
|
- for idx in pred_index:
|
|
|
- pred_texts.append(source_texts[idx])
|
|
|
- total += 1
|
|
|
- total_out_idx += sentence_bleu(
|
|
|
- [target_index],
|
|
|
- [i + 1 for i in pred_index],
|
|
|
- smoothing_function=chen_cherry.method2,
|
|
|
- )
|
|
|
- total_out_token += sentence_bleu(
|
|
|
- [" ".join(target_texts).split()],
|
|
|
- " ".join(pred_texts).split(),
|
|
|
- smoothing_function=chen_cherry.method2,
|
|
|
- )
|
|
|
-
|
|
|
- print("total: ", total)
|
|
|
- print("out_idx: ", round(100 * total_out_idx / total, 1))
|
|
|
- print("out_token: ", round(100 * total_out_token / total, 1))
|
|
|
-
|
|
|
-
|
|
|
-if __name__ == "__main__":
|
|
|
- app()
|