🚀 韓文版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許可證。