Add files using upload-large-folder tool

ed46c5d verified 3 days ago

6.6 kB

"""Shared prompt construction and model loading for Qwen3.5 recognition and translation."""
import os
import re

import torch
from PIL import Image

LANG_NAME = {"en": "English", "sv": "Swedish", "de": "German"}

OCR_PROMPT = (
    "Transcribe the runic inscription in this photograph. Output only the transliteration "
    "in lowercase Latin letters following the Rundata convention (þ, ð, æ, ø, ą, and R for "
    "the yr rune), with words separated by single spaces."
)


def mt_prompt(src, lang, rl=None):
    """rl: runic-language tag (lang_tags.py); with it the prompt carries <rl> <2lang> tags before the input."""
    tags = f"Tags: <{rl}> <2{lang}>\n" if rl else ""
    return (
        f"Translate this runic inscription, given as a Rundata-style transliteration, into "
        f"{LANG_NAME[lang]}. Output only the translation.\n\n{tags}Transliteration: {src}"
    )


MAX_IMAGE_PIXELS = int(os.environ.get("RUNIC_MAX_PIXELS", 640_000))


def load_image(path):
    im = Image.open(path).convert("RGB")
    w, h = im.size
    if w * h > MAX_IMAGE_PIXELS:
        s = (MAX_IMAGE_PIXELS / (w * h)) ** 0.5
        im = im.resize((max(32, int(w * s)), max(32, int(h * s))), Image.BICUBIC)
    return im


def user_messages(task, text=None, lang=None, rl=None):
    if task == "ocr":
        return [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": OCR_PROMPT}]}]
    return [{"role": "user", "content": [{"type": "text", "text": mt_prompt(text, lang, rl)}]}]


def _prompt_text(processor, task, text=None, lang=None, rl=None):
    msgs = user_messages(task, text, lang, rl)
    if getattr(processor, "text_only", False):
        # text-only chat templates expect the content as a plain string; some (Qwen3) silently render a list as ""
        msgs = [dict(m, content=m["content"][0]["text"]) for m in msgs]
    try:
        return processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    except TypeError:
        if task == "ocr":
            raise
        # chat templates of some text-only models (Mistral) expect the content as a plain string
        msgs = [dict(m, content=m["content"][0]["text"]) for m in msgs]
        return processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)


def prompt_text(processor, task, text=None, lang=None, rl=None):
    out = _prompt_text(processor, task, text, lang, rl)
    if task == "mt" and text and text not in out:  # a chat template that drops the user content
        raise ValueError(f"source text missing from the rendered prompt: {out[:300]!r}")
    return out


def _check_tokenizer(model_path, tok):
    """Some checkpoints (Ministral-8B-Instruct-2410 under transformers 5.x) load as a slow LlamaTokenizer that drops
    spaces when encoding; fall back to the fast tokenizer built from tokenizer.json, keeping the chat template and
    special tokens."""
    probe = "had this stone raised"
    if tok.decode(tok(probe, add_special_tokens=False)["input_ids"]).strip() == probe:
        return tok
    from transformers import PreTrainedTokenizerFast
    f = os.path.join(model_path, "tokenizer.json")
    fast = PreTrainedTokenizerFast(tokenizer_file=f, bos_token=tok.bos_token, eos_token=tok.eos_token,
                                   unk_token=tok.unk_token, pad_token=tok.pad_token)
    fast.chat_template = tok.chat_template
    if fast.decode(fast(probe, add_special_tokens=False)["input_ids"]).strip() != probe:
        raise ValueError(f"tokenizer of {model_path} does not round-trip text")
    print("tokenizer: slow tokenizer drops spaces, using tokenizer.json", flush=True)
    return fast


def load_model(model_path, adapter=None, trainable=False):
    """Vision-language checkpoints load as image-text-to-text models with their processor; text-only checkpoints
    (alternative translators) load as causal LMs, and their tokenizer stands in for the processor."""
    from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor
    from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES

    cfg = AutoConfig.from_pretrained(model_path)
    processor = AutoProcessor.from_pretrained(model_path)
    if cfg.model_type in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES:
        model = AutoModelForImageTextToText.from_pretrained(
            model_path, dtype=torch.bfloat16, attn_implementation="sdpa"
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_path, dtype=torch.bfloat16, attn_implementation="eager" if "gemma" in cfg.model_type else "sdpa"
        )
    if not hasattr(processor, "tokenizer"):
        processor = _check_tokenizer(model_path, processor)
        processor.tokenizer = processor
        processor.text_only = True
        if processor.pad_token is None:  # Mistral tokenizers have no pad token; padded positions are masked anyway
            processor.pad_token = processor.unk_token or processor.eos_token
    if adapter:
        from peft import PeftModel

        model = PeftModel.from_pretrained(model, adapter, is_trainable=trainable)
    return processor, model


def _has_token(tokenizer, t):
    i = tokenizer.convert_tokens_to_ids(t)
    return i is not None and i >= 0 and i != tokenizer.unk_token_id


def end_text(tokenizer):
    """Text that closes an assistant turn in the training target."""
    if _has_token(tokenizer, "<|im_end|>"):
        return "<|im_end|>\n"
    if _has_token(tokenizer, "<end_of_turn>"):
        return "<end_of_turn>\n"
    return tokenizer.eos_token


def encode_kwargs(processor):
    """Text-only chat templates (Gemma) already contain <bos>; do not add a second one."""
    return {"add_special_tokens": False} if getattr(processor, "text_only", False) else {}


def stop_ids(tokenizer):
    ids = {tokenizer.convert_tokens_to_ids("<|im_end|>")}
    if _has_token(tokenizer, "<end_of_turn>"):
        ids.add(tokenizer.convert_tokens_to_ids("<end_of_turn>"))
    if tokenizer.eos_token_id is not None:
        ids.add(tokenizer.eos_token_id)
    return sorted(i for i in ids if i is not None and i >= 0)


_THINK = re.compile(r"<think>.*?</think>", re.S)


def clean_generation(text):
    text = _THINK.sub("", text)
    text = text.replace("<think>", "").replace("</think>", "")
    for tok in ("<|im_end|>", "<|im_start|>", "<|endoftext|>", "<end_of_turn>", "<eos>", "</s>", "<unk>"):
        text = text.split(tok)[0]
    text = re.sub(r"^\s*(assistant|user|model)\s*\n", "", text)
    return text.strip()

	"""Shared prompt construction and model loading for Qwen3.5 recognition and translation."""
	import os
	import re

	import torch
	from PIL import Image

	LANG_NAME = {"en": "English", "sv": "Swedish", "de": "German"}

	OCR_PROMPT = (
	"Transcribe the runic inscription in this photograph. Output only the transliteration "
	"in lowercase Latin letters following the Rundata convention (þ, ð, æ, ø, ą, and R for "
	"the yr rune), with words separated by single spaces."
	)


	def mt_prompt(src, lang, rl=None):
	"""rl: runic-language tag (lang_tags.py); with it the prompt carries <rl> <2lang> tags before the input."""
	tags = f"Tags: <{rl}> <2{lang}>\n" if rl else ""
	return (
	f"Translate this runic inscription, given as a Rundata-style transliteration, into "
	f"{LANG_NAME[lang]}. Output only the translation.\n\n{tags}Transliteration: {src}"
	)


	MAX_IMAGE_PIXELS = int(os.environ.get("RUNIC_MAX_PIXELS", 640_000))


	def load_image(path):
	im = Image.open(path).convert("RGB")
	w, h = im.size
	if w * h > MAX_IMAGE_PIXELS:
	s = (MAX_IMAGE_PIXELS / (w * h)) ** 0.5
	im = im.resize((max(32, int(w * s)), max(32, int(h * s))), Image.BICUBIC)
	return im


	def user_messages(task, text=None, lang=None, rl=None):
	if task == "ocr":
	return [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": OCR_PROMPT}]}]
	return [{"role": "user", "content": [{"type": "text", "text": mt_prompt(text, lang, rl)}]}]


	def _prompt_text(processor, task, text=None, lang=None, rl=None):
	msgs = user_messages(task, text, lang, rl)
	if getattr(processor, "text_only", False):
	# text-only chat templates expect the content as a plain string; some (Qwen3) silently render a list as ""
	msgs = [dict(m, content=m["content"][0]["text"]) for m in msgs]
	try:
	return processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)
	except TypeError:
	if task == "ocr":
	raise
	# chat templates of some text-only models (Mistral) expect the content as a plain string
	msgs = [dict(m, content=m["content"][0]["text"]) for m in msgs]
	return processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)


	def prompt_text(processor, task, text=None, lang=None, rl=None):
	out = _prompt_text(processor, task, text, lang, rl)
	if task == "mt" and text and text not in out: # a chat template that drops the user content
	raise ValueError(f"source text missing from the rendered prompt: {out[:300]!r}")
	return out


	def _check_tokenizer(model_path, tok):
	"""Some checkpoints (Ministral-8B-Instruct-2410 under transformers 5.x) load as a slow LlamaTokenizer that drops
	spaces when encoding; fall back to the fast tokenizer built from tokenizer.json, keeping the chat template and
	special tokens."""
	probe = "had this stone raised"
	if tok.decode(tok(probe, add_special_tokens=False)["input_ids"]).strip() == probe:
	return tok
	from transformers import PreTrainedTokenizerFast
	f = os.path.join(model_path, "tokenizer.json")
	fast = PreTrainedTokenizerFast(tokenizer_file=f, bos_token=tok.bos_token, eos_token=tok.eos_token,
	unk_token=tok.unk_token, pad_token=tok.pad_token)
	fast.chat_template = tok.chat_template
	if fast.decode(fast(probe, add_special_tokens=False)["input_ids"]).strip() != probe:
	raise ValueError(f"tokenizer of {model_path} does not round-trip text")
	print("tokenizer: slow tokenizer drops spaces, using tokenizer.json", flush=True)
	return fast


	def load_model(model_path, adapter=None, trainable=False):
	"""Vision-language checkpoints load as image-text-to-text models with their processor; text-only checkpoints
	(alternative translators) load as causal LMs, and their tokenizer stands in for the processor."""
	from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor
	from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES

	cfg = AutoConfig.from_pretrained(model_path)
	processor = AutoProcessor.from_pretrained(model_path)
	if cfg.model_type in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES:
	model = AutoModelForImageTextToText.from_pretrained(
	model_path, dtype=torch.bfloat16, attn_implementation="sdpa"
	)
	else:
	model = AutoModelForCausalLM.from_pretrained(
	model_path, dtype=torch.bfloat16, attn_implementation="eager" if "gemma" in cfg.model_type else "sdpa"
	)
	if not hasattr(processor, "tokenizer"):
	processor = _check_tokenizer(model_path, processor)
	processor.tokenizer = processor
	processor.text_only = True
	if processor.pad_token is None: # Mistral tokenizers have no pad token; padded positions are masked anyway
	processor.pad_token = processor.unk_token or processor.eos_token
	if adapter:
	from peft import PeftModel

	model = PeftModel.from_pretrained(model, adapter, is_trainable=trainable)
	return processor, model


	def _has_token(tokenizer, t):
	i = tokenizer.convert_tokens_to_ids(t)
	return i is not None and i >= 0 and i != tokenizer.unk_token_id


	def end_text(tokenizer):
	"""Text that closes an assistant turn in the training target."""
	if _has_token(tokenizer, "<\|im_end\|>"):
	return "<\|im_end\|>\n"
	if _has_token(tokenizer, "<end_of_turn>"):
	return "<end_of_turn>\n"
	return tokenizer.eos_token


	def encode_kwargs(processor):
	"""Text-only chat templates (Gemma) already contain <bos>; do not add a second one."""
	return {"add_special_tokens": False} if getattr(processor, "text_only", False) else {}


	def stop_ids(tokenizer):
	ids = {tokenizer.convert_tokens_to_ids("<\|im_end\|>")}
	if _has_token(tokenizer, "<end_of_turn>"):
	ids.add(tokenizer.convert_tokens_to_ids("<end_of_turn>"))
	if tokenizer.eos_token_id is not None:
	ids.add(tokenizer.eos_token_id)
	return sorted(i for i in ids if i is not None and i >= 0)


	_THINK = re.compile(r"<think>.*?</think>", re.S)


	def clean_generation(text):
	text = _THINK.sub("", text)
	text = text.replace("<think>", "").replace("</think>", "")
	for tok in ("<\|im_end\|>", "<\|im_start\|>", "<\|endoftext\|>", "<end_of_turn>", "<eos>", "</s>", "<unk>"):
	text = text.split(tok)[0]
	text = re.sub(r"^\s(assistant\|user\|model)\s\n", "", text)
	return text.strip()