Implementing

[kirudang]

Hi there,

As I want to implemente this watermark, could you help invoke a sample code to integrate to model.generate mode? Like the KGW does below.

watermark_processor = WatermarkLogitsProcessor(
        vocab=list(tokenizer.get_vocab().values()),
        gamma=args.gamma,
        delta=args.delta,
        seeding_scheme=args.seeding_scheme,
        store_spike_ents=args.store_spike_ents,
        select_green_tokens=True,
        message_length=args.message_length, # specific to multi-bit watermark
        base=args.base, # specific to multi-bit watermark
        device="cuda" if (args.use_gpu and torch.cuda.is_available()) else "cpu",
        **wm_kwargs
    )

tokenized_input = tokenizer(input_text).to(model.device)
# note that if the model is on cuda, then the input is on cuda
# and thus the watermarking rng is cuda-based.
# This is a different generator than the cpu-based rng in pytorch!

output_tokens = model.generate(**tokenized_input,
                               logits_processor=LogitsProcessorList([watermark_processor]))

# if decoder only model, then we need to isolate the
# newly generated tokens as only those are watermarked, the input/prompt is not
output_tokens = output_tokens[:,tokenized_input["input_ids"].shape[-1]:]

output_text = tokenizer.batch_decode(output_without_watermark, skip_special_tokens=True)[0]

I saw your generation_pipeline but can i skip data_collator and generation_partial?

[bangawayoo]

Hi, Have you tried simply replacing WatermarkLogitsProcessor with our processor from mb_watermark_processor? This would mean importing the following line

from mb_watermark_processor import WatermarkLogitsProcessor
# from watermark_processor import WatermarkLogitsProcessor

If you have any more questions, please open an issue in this repository.