diff --git a/README.md b/README.md index 88f33bb..0d31c4f 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@
English | 中文版本 @@ -23,7 +23,7 @@ python=3.10 pytorch -> Note: CUDA version >= 12.0 have been fully tested. +> Note: Only CUDA version >= 12.0 have been fully tested, so we recommend using CUDA version>=12.0 ## Getting Started @@ -33,7 +33,7 @@ pytorch git clone https://github.com/OleehyO/TexTeller ``` -2. After [pytorch installation](https://pytorch.org/get-started), install the required packages: +2. After [pytorch installation](https://pytorch.org/get-started/locally/#start-locally), install the required packages: ```bash pip install -r requirements.txt @@ -47,7 +47,7 @@ pytorch #+e.g. python inference.py -img "./img.jpg" -cuda ``` - > chekpoints will be downloaded in your first run. + > Checkpoints will be downloaded in your first run. ## Web Demo @@ -57,9 +57,9 @@ You can also run the web demo by navigating to the `TexTeller/src` directory and ./start_web.sh ``` -Then go to `http://localhost:8000` in your browser to run TexTeller in the web. +Then go to `http://localhost:8501` in your browser to run TexTeller in the web. -> You can change the default settings in start_web.sh, such as inference with GPU(e.g. `USE_CUDA=True`) or increase the number of beams(e.g. `NUM_BEAM=3`) for higher accuracy. +> You can change the default settings in `start_web.sh`, such as inference with GPU(e.g. `USE_CUDA=True`) or increase the number of beams(e.g. `NUM_BEAM=3`) for higher accuracy. ## API @@ -82,13 +82,13 @@ You can pass the following arguments to the `serve.py` script to get custom infe | `--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) | -> Client demo can be found in `TexTeller/client/demo.py`. +> Client demo can be found in `TexTeller/client/demo.py`, you can refer to `demo.py` to send requests to the server. ## Training ### Dataset -We provide a dataset example in `TexTeller/src/models/ocr_model/train/dataset`, and you can place your own images in the `images` directory and annotate the corresponding formula for each image in `formulas.jsonl` +We provide a dataset example in `TexTeller/src/models/ocr_model/train/dataset`, and you can place your own images in the `images` directory and annotate the corresponding formula for each image in `formulas.jsonl`. After the dataset is ready, you should **change the `DIR_URL` variable** in `.../dataset/loader.py` to the path of your dataset. @@ -96,7 +96,7 @@ After the dataset is ready, you should **change the `DIR_URL` variable** in `... If you are using a different dataset, you may need to retrain the tokenizer to match your specific vocabulary. After setting up the dataset, you can do this by: -1. Change the line `new_tokenizer.save_pretrained('./your_dir_name')` in `TexTeller/src/models/ocr_model/tokenizer/train.py` to your desired output directory name.` +1. Change the line `new_tokenizer.save_pretrained('./your_dir_name')` in `TexTeller/src/models/ocr_model/tokenizer/train.py` to your desired output directory name. > To use a different vocabulary size, you should modify the `VOCAB_SIZE` parameter in the `TexTeller/src/models/globals.py`. 2. Running the following command **under `TexTeller/src` directory**: @@ -114,6 +114,7 @@ python -m models.ocr_model.train.train ``` You can set your own tokenizer and checkpoint path(or fine-tune the default model checkpoint if you don't use your own tokenizer while keeping the same model architecture) in `TexTeller/src/models/ocr_model/train/train.py`. +> Please refer to `train.py` for more details. Model architecture and training hyperparameters can be adjusted in `TexTeller/src/globals.py` and `TexTeller/src/models/ocr_model/train/train_args.py`. diff --git a/assets/README_zh.md b/assets/README_zh.md new file mode 100644 index 0000000..d8a9e59 --- /dev/null +++ b/assets/README_zh.md @@ -0,0 +1,162 @@ +
+ +TexTeller是一个基于ViT的端到端公式识别模型,可以把图片转换为对应的latex公式 + +TexTeller用了550K的图片-公式对进行训练(数据集可以在[这里](https://huggingface.co/datasets/OleehyO/latex-formulas)获取),相比于[LaTeX-OCR](https://github.com/lukas-blecher/LaTeX-OCR)(使用了一个100K的数据集),TexTeller具有更强的泛化能力以及更高的精确度,可以覆盖你大部分的使用场景。 + +> 我们马上就会发布一个使用5.5M数据集进行训练的TexTeller checkpoint + +## 前置条件 + +python=3.10 + +pytorch + +> 注意: 只有CUDA版本>= 12.0被完全测试过,所以最好使用>= 12.0的CUDA版本 + + +## Getting Started + +1. 克隆本仓库: + + ```bash + git clone https://github.com/OleehyO/TexTeller + ``` + +2. [安装pytorch](https://pytorch.org/get-started/locally/#start-locally)后,再安装本项目的依赖包: + + ```bash + pip install -r requirements.txt + ``` + +3. 进入`TexTeller/src`目录,在终端运行以下命令进行推理: + + ```bash + 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包 + + ```bash + pip install -U "huggingface_hub[cli]" + ``` + +2. 在能连接Hugging Face的机器上下载模型权重: + + ```bash + 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脚本 + + ```bash + 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')` + +## Web Demo + +要想启动web demo,你需要先进入 `TexTeller/src` 目录,然后运行以下命令 + +```bash +./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](https://github.com/ray-project/ray)来对外提供一个TexTeller的API接口,通过使用这个接口,你可以把TexTeller整合到自己的项目里。要想启动server,你需要先进入`TexTeller/src`目录然后运行以下命令: + +```bash +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/ocr_model/tokenizer/train.py`中,修改`new_tokenizer.save_pretrained('./your_dir_name')`为你自定义的输出目录 + > 如果要用一个不一样大小的字典(默认1W个token),你需要在 `TexTeller/src/models/globals.py`中修改`VOCAB_SIZE`变量 + +2. **在 `TexTeller/src` 目录下**运行以下命令: + + ```bash + python -m models.ocr_model.tokenizer.train + ``` + +### Train the model + +要想训练模型, 你需要在`TexTeller/src`目录下运行以下命令: + +```bash +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](https://github.com/huggingface/transformers)库, 所以你可以参考他们提供的[文档](https://huggingface.co/docs/transformers/v4.32.1/main_classes/trainer#transformers.TrainingArguments)来获取更多训练参数的细节以及配置。 + +## To-Do + +- [ ] 使用更大的数据集来训练模型(5.5M样本,即将发布) + +- [ ] 推理加速 + +- [ ] ... + +## Acknowledgements + +Thanks to [LaTeX-OCR](https://github.com/lukas-blecher/LaTeX-OCR) which has brought me a lot of inspiration, and [im2latex-100K](https://zenodo.org/records/56198#.V2px0jXT6eA) which enriches our dataset. diff --git a/src/models/ocr_model/model/TexTeller.py b/src/models/ocr_model/model/TexTeller.py index 740ccc5..4b059bd 100644 --- a/src/models/ocr_model/model/TexTeller.py +++ b/src/models/ocr_model/model/TexTeller.py @@ -30,14 +30,14 @@ class TexTeller(VisionEncoderDecoderModel): @classmethod def from_pretrained(cls, model_path: str = None): - if model_path is None or model_path == cls.REPO_NAME: + if model_path is None or model_path == 'default': return VisionEncoderDecoderModel.from_pretrained(cls.REPO_NAME) model_path = Path(model_path).resolve() return VisionEncoderDecoderModel.from_pretrained(str(model_path)) @classmethod def get_tokenizer(cls, tokenizer_path: str = None) -> RobertaTokenizerFast: - if tokenizer_path is None or tokenizer_path == cls.REPO_NAME: + if tokenizer_path is None or tokenizer_path == 'default': return RobertaTokenizerFast.from_pretrained(cls.REPO_NAME) tokenizer_path = Path(tokenizer_path).resolve() return RobertaTokenizerFast.from_pretrained(str(tokenizer_path)) diff --git a/src/start_web.sh b/src/start_web.sh index 7d32eed..450dff2 100755 --- a/src/start_web.sh +++ b/src/start_web.sh @@ -1,9 +1,9 @@ #!/usr/bin/env bash set -exu -export CHECKPOINT_DIR="OleehyO/TexTeller" -export TOKENIZER_DIR="OleehyO/TexTeller" -export USE_CUDA=True # True or False (case-sensitive) -export NUM_BEAM=5 +export CHECKPOINT_DIR="default" +export TOKENIZER_DIR="default" +export USE_CUDA=False # True or False (case-sensitive) +export NUM_BEAM=1 streamlit run web.py