stevafernandes/Fine-Tuning-HunyuanOCR

1"""MuseumSCAT specimen label transcription — Hugging Face Space.2 3Loads HunyuanOCR-1.5 with an optional fine-tuned LoRA adapter and extracts4verbatimDate and verbatimLocality (with confidence scores) from a specimen5label image.6 7Adapter resolution order:8  1. ADAPTER_REPO environment variable (Hugging Face model repo id)9  2. ./adapter directory bundled with the Space10  3. none — base model11"""12 13import json14import os15import re16 17import torch18from PIL import Image19from transformers import AutoProcessor, HunYuanVLForConditionalGeneration20 21from ui import build_demo22 23MODEL_ID = "tencent/HunyuanOCR"24MISSING = "MISSING"25# Longest image side. Must match the resolution the adapter was trained at26# (prepare_data.py --max-side).27MAX_SIDE = 153628 29PROMPT = (30    "Read the specimen labels in this image. Transcribe the collection date and the "31    "collection locality exactly as written on the labels, using the Danish alphabet "32    '(replace historical umlauts). Answer with JSON only, in the form '33    '{"verbatimDate": "...", "verbatimLocality": "..."}. '34    'Use "MISSING" for a field that is not present on any label.'35)36 37try:38    import spaces  # ZeroGPU support; absent outside Hugging Face Spaces39 40    gpu_decorator = spaces.GPU41except ImportError:42 43    def gpu_decorator(fn):44        return fn45 46 47def resolve_adapter() -> str | None:48    repo = os.environ.get("ADAPTER_REPO", "").strip()49    if repo:50        return repo51    local = os.path.join(os.path.dirname(os.path.abspath(__file__)), "adapter")52    if os.path.exists(os.path.join(local, "adapter_config.json")):53        return local54    return None55 56 57def pick_device() -> str:58    if torch.cuda.is_available():59        return "cuda"60    if torch.backends.mps.is_available():61        return "mps"62    return "cpu"63 64 65DEVICE = pick_device()66ADAPTER = resolve_adapter()67 68processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=False)69model = HunYuanVLForConditionalGeneration.from_pretrained(70    MODEL_ID,71    dtype=torch.bfloat16 if DEVICE != "cpu" else torch.float32,72    attn_implementation="sdpa",73)74if ADAPTER:75    from peft import PeftModel76 77    model = PeftModel.from_pretrained(model, ADAPTER)78    model = model.merge_and_unload()79model = model.to(DEVICE)80model.eval()81 82 83def parse_prediction(text: str) -> dict:84    text = text.strip()85    decoder = json.JSONDecoder()86    for match in re.finditer(r"\{", text):87        try:88            obj, _ = decoder.raw_decode(text, match.start())89        except ValueError:90            continue91        if isinstance(obj, dict):92            date = obj.get("verbatimDate")93            locality = obj.get("verbatimLocality")94            return {95                "verbatimDate": str(date) if date is not None else MISSING,96                "verbatimLocality": str(locality) if locality is not None else MISSING,97            }98    return {"verbatimDate": MISSING, "verbatimLocality": MISSING}99 100 101def field_confidence(field_name: str, field_value: str, token_ids, logprobs, tokenizer) -> float:102    if not logprobs:103        return 0.0104    overall = float(torch.exp(torch.tensor(logprobs).mean()))105    if not field_value:106        return overall107    decoded = tokenizer.decode(token_ids, skip_special_tokens=True)108    key_pos = decoded.find(f'"{field_name}"')109    start = decoded.find(field_value, key_pos + len(field_name) + 2) if key_pos != -1 else -1110    if start == -1:111        start = decoded.find(field_value)112    if start == -1:113        return overall114    end = start + len(field_value)115    span_logprobs = []116    prev_len = 0117    for i in range(len(token_ids)):118        cur_len = len(tokenizer.decode(token_ids[: i + 1], skip_special_tokens=True).rstrip("�"))119        if cur_len > start and prev_len < end and i < len(logprobs):120            span_logprobs.append(logprobs[i])121        prev_len = cur_len122        if prev_len >= end:123            break124    if not span_logprobs:125        return overall126    return float(torch.exp(torch.tensor(span_logprobs).mean()))127 128 129@gpu_decorator130def transcribe(image: Image.Image):131    if image is None:132        return "", "", "", "", ""133    image = image.convert("RGB")134    if max(image.size) > MAX_SIDE:135        image.thumbnail((MAX_SIDE, MAX_SIDE), Image.LANCZOS)136 137    messages = [138        {"role": "system", "content": ""},139        {140            "role": "user",141            "content": [142                {"type": "image"},143                {"type": "text", "text": PROMPT},144            ],145        },146    ]147    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)148    inputs = processor(text=[text], images=image, return_tensors="pt").to(DEVICE)149    prompt_len = inputs["input_ids"].shape[1]150 151    with torch.inference_mode():152        out = model.generate(153            **inputs,154            max_new_tokens=128,155            do_sample=False,156            repetition_penalty=1.08,157            pad_token_id=processor.tokenizer.eos_token_id,158            output_scores=True,159            return_dict_in_generate=True,160        )161    token_ids = out.sequences[0][prompt_len:]162    logprobs = []163    for i, scores in enumerate(out.scores):164        if i >= len(token_ids):165            break166        step_logprobs = torch.log_softmax(scores[0].float(), dim=-1)167        logprobs.append(step_logprobs[token_ids[i]].item())168 169    decoded = processor.tokenizer.decode(token_ids, skip_special_tokens=True)170    parsed = parse_prediction(decoded)171    token_list = token_ids.tolist()172    tokenizer = processor.tokenizer173    date_conf = field_confidence("verbatimDate", parsed["verbatimDate"], token_list, logprobs, tokenizer)174    loc_conf = field_confidence("verbatimLocality", parsed["verbatimLocality"], token_list, logprobs, tokenizer)175 176    return (177        parsed["verbatimDate"],178        f"{date_conf:.3f}",179        parsed["verbatimLocality"],180        f"{loc_conf:.3f}",181        decoded,182    )183 184 185demo = build_demo(transcribe)186 187if __name__ == "__main__":188    demo.launch()189