Delete benchmark.py

benchmark.py

@@ -1,226 +0,0 @@
-import argparse
-import os.path
-import random
-import time
-from functools import partial
-
-import evaluate
-from tabulate import tabulate
-from tqdm import tqdm
-
-from texify.inference import batch_inference
-from texify.model.model import load_model
-from texify.model.processor import load_processor
-from PIL import Image
-from texify.settings import settings
-import json
-import base64
-import io
-from rapidfuzz.distance import Levenshtein
-
-
-def normalize_text(text):
-    # Replace fences
-    text = text.replace("$", "")
-    text = text.replace("\[", "")
-    text = text.replace("\]", "")
-    text = text.replace("\(", "")
-    text = text.replace("\)", "")
-    text = text.strip()
-    return text
-
-
-def score_text(predictions, references):
-    bleu = evaluate.load("bleu")
-    bleu_results = bleu.compute(predictions=predictions, references=references)
-
-    meteor = evaluate.load('meteor')
-    meteor_results = meteor.compute(predictions=predictions, references=references)
-
-    lev_dist = []
-    for p, r in zip(predictions, references):
-        lev_dist.append(Levenshtein.normalized_distance(p, r))
-
-    return {
-        'bleu': bleu_results["bleu"],
-        'meteor': meteor_results['meteor'],
-        'edit': sum(lev_dist) / len(lev_dist)
-    }
-
-
-def image_to_pil(image):
-    decoded = base64.b64decode(image)
-    return Image.open(io.BytesIO(decoded))
-
-
-def load_images(source_data):
-    images = [sd["image"] for sd in source_data]
-    images = [image_to_pil(image) for image in images]
-    return images
-
-
-def inference_texify(source_data, model, processor):
-    images = load_images(source_data)
-
-    write_data = []
-    for i in tqdm(range(0, len(images), settings.BATCH_SIZE), desc="Texify inference"):
-        batch = images[i:i+settings.BATCH_SIZE]
-        text = batch_inference(batch, model, processor)
-        for j, t in enumerate(text):
-            eq_idx = i + j
-            write_data.append({"text": t, "equation": source_data[eq_idx]["equation"]})
-
-    return write_data
-
-
-def inference_pix2tex(source_data):
-    from pix2tex.cli import LatexOCR
-    model = LatexOCR()
-
-    images = load_images(source_data)
-    write_data = []
-    for i in tqdm(range(len(images)), desc="Pix2tex inference"):
-        try:
-            text = model(images[i])
-        except ValueError:
-            # Happens when resize fails
-            text = ""
-        write_data.append({"text": text, "equation": source_data[i]["equation"]})
-
-    return write_data
-
-
-def image_to_bmp(image):
-    img_out = io.BytesIO()
-    image.save(img_out, format="BMP")
-    return img_out
-
-
-def inference_nougat(source_data, batch_size=1):
-    import torch
-    from nougat.postprocessing import markdown_compatible
-    from nougat.utils.checkpoint import get_checkpoint
-    from nougat.utils.dataset import ImageDataset
-    from nougat.utils.device import move_to_device
-    from nougat import NougatModel
-
-    # Load images, then convert to bmp format for nougat
-    images = load_images(source_data)
-    images = [image_to_bmp(image) for image in images]
-    predictions = []
-
-    ckpt = get_checkpoint(None, model_tag="0.1.0-small")
-    model = NougatModel.from_pretrained(ckpt)
-    if settings.TORCH_DEVICE_MODEL != "cpu":
-        move_to_device(model, bf16=settings.CUDA, cuda=settings.CUDA)
-    model.eval()
-
-    dataset = ImageDataset(
-        images,
-        partial(model.encoder.prepare_input, random_padding=False),
-    )
-
-    # Batch sizes higher than 1 explode memory usage on CPU/MPS
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=batch_size,
-        pin_memory=True,
-        shuffle=False,
-    )
-
-    for idx, sample in tqdm(enumerate(dataloader), desc="Nougat inference", total=len(dataloader)):
-        model.config.max_length = settings.MAX_TOKENS
-        model_output = model.inference(image_tensors=sample, early_stopping=False)
-        output = [markdown_compatible(o) for o in model_output["predictions"]]
-        predictions.extend(output)
-    return predictions
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Benchmark the performance of texify.")
-    parser.add_argument("--data_path", type=str, help="Path to JSON file with source images/equations", default=os.path.join(settings.DATA_DIR, "bench_data.json"))
-    parser.add_argument("--result_path", type=str, help="Path to JSON file to save results to.", default=os.path.join(settings.DATA_DIR, "bench_results.json"))
-    parser.add_argument("--max", type=int, help="Maximum number of images to benchmark.", default=None)
-    parser.add_argument("--pix2tex", action="store_true", help="Run pix2tex scoring", default=False)
-    parser.add_argument("--nougat", action="store_true", help="Run nougat scoring", default=False)
-    args = parser.parse_args()
-
-    source_path = os.path.abspath(args.data_path)
-    result_path = os.path.abspath(args.result_path)
-    os.makedirs(os.path.dirname(result_path), exist_ok=True)
-    model = load_model()
-    processor = load_processor()
-
-    with open(source_path, "r") as f:
-        source_data = json.load(f)
-
-    if args.max:
-        random.seed(1)
-        source_data = random.sample(source_data, args.max)
-
-    start = time.time()
-    predictions = inference_texify(source_data, model, processor)
-    times = {"texify": time.time() - start}
-    text = [normalize_text(p["text"]) for p in predictions]
-    references = [normalize_text(p["equation"]) for p in predictions]
-
-    scores = score_text(text, references)
-
-    write_data = {
-        "texify": {
-            "scores": scores,
-            "text": [{"prediction": p, "reference": r} for p, r in zip(text, references)]
-        }
-    }
-
-    if args.pix2tex:
-        start = time.time()
-        predictions = inference_pix2tex(source_data)
-        times["pix2tex"] = time.time() - start
-
-        p_text = [normalize_text(p["text"]) for p in predictions]
-
-        p_scores = score_text(p_text, references)
-
-        write_data["pix2tex"] = {
-            "scores": p_scores,
-            "text": [{"prediction": p, "reference": r} for p, r in zip(p_text, references)]
-        }
-
-    if args.nougat:
-        start = time.time()
-        predictions = inference_nougat(source_data)
-        times["nougat"] = time.time() - start
-        n_text = [normalize_text(p) for p in predictions]
-
-        n_scores = score_text(n_text, references)
-
-        write_data["nougat"] = {
-            "scores": n_scores,
-            "text": [{"prediction": p, "reference": r} for p, r in zip(n_text, references)]
-        }
-
-    score_table = []
-    score_headers = ["bleu", "meteor", "edit"]
-    score_dirs = ["⬆", "⬆", "⬇", "⬇"]
-
-    for method in write_data.keys():
-        score_table.append([method, *[write_data[method]["scores"][h] for h in score_headers], times[method]])
-
-    score_headers.append("time taken (s)")
-    score_headers = [f"{h} {d}" for h, d in zip(score_headers, score_dirs)]
-    print()
-    print(tabulate(score_table, headers=["Method", *score_headers]))
-    print()
-    print("Higher is better for BLEU and METEOR, lower is better for edit distance and time taken.")
-    print("Note that pix2tex is unbatched (I couldn't find a batch inference method in the docs), so time taken is higher than it should be.")
-
-    with open(result_path, "w") as f:
-        json.dump(write_data, f, indent=4)
-
-
-if __name__ == "__main__":
-    main()
-
-
-