sign/WeLT-string-repetition

from pathlib import Path

import torch
from transformers import GenerationConfig
from transformers.trainer_utils import get_last_checkpoint

from welt.model import WordLatentTransformerForCausalLM
from welt.processor import TextImageProcessor


with torch.no_grad():
    model = WordLatentTransformerForCausalLM.from_pretrained("sign/WeLT-string-repetition")
    processor = TextImageProcessor.from_pretrained(model_path)
    model.eval()

    texts = [
        # Texts from validation set
        "<text>\x0EWouldn't it be more cruel for society to let people die... - ... when with some effort it could save them?\x0F<repeat> ",
        "<text>\x0ESuperman's exact opposite who lives in the backwards Bizarro World.\x0F<repeat> ",
        "<text>\x0EYOu dOn't know the half Of it.\x0F<repeat> ",
    ]

    inputs = processor(texts, collated=True, packed=False)

    outputs = model.generate(
        **inputs,
        processor=processor,
        max_generated_words=32,
    )
    for text, output in zip(texts, outputs, strict=False):
        print(f"Generated for '{text}': {output}")