Ko Trocr Base Nsmc News Chatbot

🚀 韩文版TrOCR模型（概念验证）

本项目为韩文语言训练了TrOCR模型用于概念验证。由于TrOCR尚未发布包含韩语的多语言模型，我们通过收集特定数据集进行训练。基于此模型，建议收集更多数据进行第一阶段的额外训练或第二阶段的微调。

🚀 快速开始

推理示例

from transformers import TrOCRProcessor, VisionEncoderDecoderModel, AutoTokenizer
import requests 
from io import BytesIO
from PIL import Image

processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") 
model = VisionEncoderDecoderModel.from_pretrained("daekeun-ml/ko-trocr-base-nsmc-news-chatbot")
tokenizer = AutoTokenizer.from_pretrained("daekeun-ml/ko-trocr-base-nsmc-news-chatbot")

url = "https://raw.githubusercontent.com/aws-samples/sm-kornlp/main/trocr/sample_imgs/news_1.jpg"
response = requests.get(url)
img = Image.open(BytesIO(response.content))

pixel_values = processor(img, return_tensors="pt").pixel_values 
generated_ids = model.generate(pixel_values, max_length=64)
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] 
print(generated_text)

✨ 主要特性

• 针对韩语进行训练，为韩语的图像文本识别提供了概念验证模型。
• 基于TrOCR架构，使用了特定的编码器和解码器模型进行训练。

📦 安装指南

文档未提及安装相关内容，可参考作者的Github仓库代码进行环境配置。

📚 详细文档

数据收集

文本数据

我们通过处理三种类型的数据集来创建训练数据：

• 新闻摘要数据集：https://huggingface.co/datasets/daekeun-ml/naver-news-summarization-ko
• Naver电影情感分类数据集：https://github.com/e9t/nsmc
• 聊天机器人数据集：https://github.com/songys/Chatbot_data

为了高效收集数据，使用句子分隔库（Kiwi Python包装器；https://github.com/bab2min/kiwipiepy）对每个句子进行分隔，最终收集了637,401个样本。

图像数据

图像数据使用TrOCR论文中介绍的TextRecognitionDataGenerator（https://github.com/Belval/TextRecognitionDataGenerator）生成。以下是生成图像的代码片段：

python3 ./trdg/run.py -i ocr_dataset_poc.txt -w 5 -t {num_cores} -f 64 -l ko -c {num_samples} -na 2 --output_dir {dataset_dir}

训练

基础模型

编码器模型使用了facebook/deit-base-distilled-patch16-384，解码器模型使用了klue/roberta-base。这比从microsoft/trocr-base-stage1开始训练权重更容易。

参数

我们使用了启发式参数，没有进行单独的超参数调整：

• 学习率 = 4e-5
• 训练轮数 = 25
• 混合精度训练（fp16） = True
• 最大长度 = 64

模型使用

所有数据收集和模型训练所需的代码已发布在作者的Github上：

• https://github.com/daekeun-ml/sm-kornlp-usecases/tree/main/trocr

🔧 技术细节

本项目基于TrOCR架构，针对韩语进行了特定的训练。在数据收集阶段，通过处理多种文本数据集并使用工具生成图像数据。训练时选择了合适的编码器和解码器模型，并使用启发式参数进行训练，为韩语的图像文本识别提供了一种可行的解决方案。

📄 许可证

本项目使用MIT许可证。