Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2from peft import PeftModel
3import torch
4from PIL import Image
5import os
6from tqdm import tqdm
7import json
8
9# Configuration
10device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
11model_name = "openbmb/MiniCPM-Llama3-V-2_5"
12abbr_adapters = "magistermilitum/Tridis_HTR_MiniCPM_ABBR"
13not_abbr_adapters = "magistermilitum/Tridis_HTR_MiniCPM"
14
15image_folder = "/your/images/folder/path"
16
17class TranscriptionModel:
18 """Handles model loading, adapter switching, and transcription generation."""
19 def __init__(self, model_name, abbr_adapters, not_abbr_adapters, device):
20 self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
21 self.base_model = AutoModelForCausalLM.from_pretrained(
22 model_name, trust_remote_code=True, attn_implementation='sdpa', torch_dtype=torch.bfloat16, token=True
23 )
24 self.base_model = PeftModel.from_pretrained(self.base_model, abbr_adapters, adapter_name="ABBR")
25 self.base_model.load_adapter(not_abbr_adapters, adapter_name="NOT_ABBR")
26 self.base_model.set_adapter("ABBR") # Set default adapter
27 self.base_model.to(device).eval()
28
29 def generate(self, adapter, image):
30 """Generate transcription for the given adapter and image."""
31 if hasattr(self.base_model, "past_key_values"):
32 self.base_model.past_key_values = None
33 self.base_model.set_adapter(adapter)
34 msgs = [{"role": "user", "content": [f"Transcribe this manuscript line in mode <{adapter}>:", image]}]
35 with torch.no_grad():
36 res = self.base_model.chat(image=image, msgs=msgs, tokenizer=self.tokenizer, max_new_tokens=128)
37 # Remove <ABBR> and <NOT_ABBR> tokens from the output
38 res = res.replace(f"<{adapter}>", "").replace(f"</{adapter}>", "")
39 return res
40
41
42class TranscriptionPipeline:
43 """Handles image processing, transcription, and result saving."""
44 def __init__(self, model, image_folder):
45 self.model = model
46 self.image_folder = image_folder
47
48 def run_inference(self):
49 """Process all images in the folder and generate transcriptions."""
50 results = []
51 for image_file in tqdm([f for f in os.listdir(self.image_folder)[:20] if f.endswith(('.png', '.jpg', '.jpeg'))]):
52 image = Image.open(os.path.join(self.image_folder, image_file)).convert("RGB")
53 print(f"\nProcessing image: {image_file}")
54
55 # Generate transcriptions for both adapters
56 transcriptions = {
57 adapter: self.model.generate(adapter, image)
58 for adapter in ["ABBR", "NOT_ABBR"]
59 }
60 for adapter, res in transcriptions.items():
61 print(f"Mode ({adapter}): {res}")
62 results.append({"image": image_file, "transcriptions": transcriptions})
63
64 #image.show() #Optional
65
66 # Save results to a JSON file
67 with open("transcriptions_results.json", "w", encoding="utf-8") as f:
68 json.dump(results, f, ensure_ascii=False, indent=4)
69
70
71# Initialize and run the pipeline
72model = TranscriptionModel(model_name, abbr_adapters, not_abbr_adapters, device)
73TranscriptionPipeline(model, image_folder).run_inference()1@misc{torres_aguilar:hal-04983305,
2 title={Dual-Style Transcription of Historical Manuscripts based on Multimodal Small Language Models with Switchable Adapters},
3 author={Torres Aguilar, Sergio},
4 url={https://hal.science/hal-04983305},
5 year={2025},
6 note = {working paper or preprint}
7}