[refactor] Init

texteller/cli/commands/inference.py

@@ -0,0 +1,51 @@
+"""
+CLI command for formula inference from images.
+"""
+
+import click
+
+from texteller.api import img2latex, load_model, load_tokenizer
+
+
+@click.command()
+@click.argument("image_path", type=click.Path(exists=True, file_okay=True, dir_okay=False))
+@click.option(
+    "--model-path",
+    type=click.Path(exists=True, file_okay=False, dir_okay=True),
+    default=None,
+    help="Path to the model dir path, if not provided, will use model from huggingface repo",
+)
+@click.option(
+    "--tokenizer-path",
+    type=click.Path(exists=True, file_okay=False, dir_okay=True),
+    default=None,
+    help="Path to the tokenizer dir path, if not provided, will use tokenizer from huggingface repo",
+)
+@click.option(
+    "--output-format",
+    type=click.Choice(["latex", "katex"]),
+    default="katex",
+    help="Output format, either latex or katex",
+)
+@click.option(
+    "--keep-style",
+    is_flag=True,
+    default=False,
+    help="Whether to keep the style of the LaTeX (e.g. bold, italic, etc.)",
+)
+def inference(image_path, model_path, tokenizer_path, output_format, keep_style):
+    """
+    CLI command for formula inference from images.
+    """
+    model = load_model(model_dir=model_path)
+    tknz = load_tokenizer(tokenizer_dir=tokenizer_path)
+
+    pred = img2latex(
+        model=model,
+        tokenizer=tknz,
+        images=[image_path],
+        out_format=output_format,
+        keep_style=keep_style,
+    )[0]
+
+    click.echo(f"Predicted LaTeX: ```\n{pred}\n```")