Skip to content

Latest commit

 

History

History
162 lines (100 loc) · 6.62 KB

README_zh.md

File metadata and controls

162 lines (100 loc) · 6.62 KB

𝚃𝚎𝚡𝚃𝚎𝚕𝚕𝚎𝚛

English | 中文版本

TexTeller_demo

TexTeller是一个基于ViT的端到端公式识别模型，可以把图片转换为对应的latex公式

TexTeller用了550K的图片-公式对进行训练(数据集可以在这里获取)，相比于LaTeX-OCR(使用了一个100K的数据集)，TexTeller具有更强的泛化能力以及更高的准确率，可以覆盖大部分的使用场景。

我们马上就会发布一个使用5.5M数据集进行训练的TexTeller checkpoint

前置条件

python=3.10

pytorch

注意: 只有CUDA版本>= 12.0被完全测试过，所以最好使用>= 12.0的CUDA版本

Getting Started

  1. 克隆本仓库:

    git clone https://github.com/OleehyO/TexTeller

  2. 安装pytorch后，再安装本项目的依赖包:

    pip install -r requirements.txt

  3. 进入TexTeller/src目录，在终端运行以下命令进行推理:

    python inference.py -img "/path/to/image.{jpg,png}" 
# use -cuda option to enable GPU inference
#+e.g. python inference.py -img "./img.jpg" -cuda

    第一次运行时会在hugging face上下载所需要的checkpoints

FAQ：无法连接到Hugging Face

默认情况下，会在Hugging Face中下载模型权重，如果你的远端服务器无法连接到Hugging Face，你可以通过以下命令进行加载：

  1. 安装huggingface hub包

    pip install -U "huggingface_hub[cli]"

  2. 在能连接Hugging Face的机器上下载模型权重:

    huggingface-cli download OleehyO/TexTeller --include "*.json" "*.bin" "*.txt" --repo-type model --local-dir "your/dir/path"

  3. 把包含权重的目录上传远端服务器，然后把TexTeller/src/models/ocr_model/model/TexTeller.py中的REPO_NAME = 'OleehyO/TexTeller'修改为REPO_NAME = 'your/dir/path'

如果你还想在训练模型时开启evaluate，你需要提前下载metric脚本并上传远端服务器：

  1. 在能连接Hugging Face的机器上下载metric脚本

    huggingface-cli download evaluate-metric/google_bleu --repo-type space --local-dir "your/dir/path"

  2. 把这个目录上传远端服务器，并在TexTeller/src/models/ocr_model/utils/metrics.py中把evaluate.load('google_bleu')改为evaluate.load('your/dir/path/google_bleu.py')

Web Demo

要想启动web demo，你需要先进入 TexTeller/src 目录，然后运行以下命令

./start_web.sh

然后在浏览器里输入http://localhost:8501就可以看到web demo

你可以改变start_web.sh的默认配置， 例如使用GPU进行推理(e.g. USE_CUDA=True) 或者增加beams的数量(e.g. NUM_BEAM=3)来获得更高的精确度

API

我们使用ray serve来对外提供一个TexTeller的API接口，通过使用这个接口，你可以把TexTeller整合到自己的项目里。要想启动server，你需要先进入TexTeller/src目录然后运行以下命令:

python serve.py  # default settings

你可以给serve.py传递以下参数来改变server的推理设置(e.g. python serve.py --use_gpu 来启动GPU推理):

Argument Description
-ckpt Path to the checkpoint file to load, default is TexTeller pretrained model.
-tknz Path to the tokenizer, default is TexTeller tokenizer.
-port Port number to run the server on, default is 8000.
--use_gpu Whether to use GPU for inference.
--num_beams Number of beams to use for beam search decoding, default is 1.
--num_replicas Number of replicas to run the server on, default is 1. You can use this to get higher throughput.
--ncpu_per_replica Number of CPU cores to use per replica, default is 1.
--ngpu_per_replica Number of GPUs to use per replica, default is 1. You can set this to 0~1 to run multiple replicas on a single GPU(if --num_replicas 2, --ngpu_per_replica 0.7, then 2 gpus are required)

一个客户端demo可以在TexTeller/client/demo.py找到，你可以参考demo.py来给server发送请求

Training

Dataset

我们在TexTeller/src/models/ocr_model/train/dataset目录中提供了一个数据集的例子，你可以把自己的图片放在images目录然后在formulas.jsonl中为每张图片标注对应的公式。

准备好数据集后，你需要在.../dataset/loader.py中把 DIR_URL变量改成你自己数据集的路径

Retrain the tokenizer

如果你使用了不一样的数据集，你可能需要重新训练tokenizer来得到一个不一样的字典。配置好数据集后，可以通过以下命令来训练自己的tokenizer：

  1. 在TexTeller/src/models/tokenizer/train.py中，修改new_tokenizer.save_pretrained('./your_dir_name')为你自定义的输出目录

    如果要用一个不一样大小的字典(默认1W个token)，你需要在 TexTeller/src/models/globals.py中修改VOCAB_SIZE变量

  2. 在 TexTeller/src 目录下运行以下命令:

    python -m models.tokenizer.train

Train the model

要想训练模型, 你需要在TexTeller/src目录下运行以下命令：

python -m models.ocr_model.train.train

你可以在TexTeller/src/models/ocr_model/train/train.py中设置自己的tokenizer和checkpoint路径（请参考train.py）。如果你使用了与TexTeller一样的架构和相同的字典，你还可以用自己的数据集来微调TexTeller的默认权重。

在TexTeller/src/globals.py和TexTeller/src/models/ocr_model/train/train_args.py中，你可以改变模型的架构以及训练的超参数。

我们的训练脚本使用了Hugging Face Transformers库, 所以你可以参考他们提供的文档来获取更多训练参数的细节以及配置。

To-Do

  • 使用更大的数据集来训练模型(5.5M样本，即将发布)

  • 推理加速

  • ...

Acknowledgements

Thanks to LaTeX-OCR which has brought me a lot of inspiration, and im2latex-100K which enriches our dataset.