Not able to run the Taskrunner or training code

Hi! Great work! Congratulations! Thanks for releasing the code!

However, I am not able to reproduce the results for taskrunners using any of the allenai/uio2-large, allenai/uio2-xl or allenai/uio2-xxl models. I am using python 3.8 with transformers 4.38.2. Also downloaded the tokenizer.model from the official llama repo.

I tried using the code as:

import torch
from uio2.model import UnifiedIOModel
from uio2.preprocessing import UnifiedIOPreprocessor
from uio2.runner import TaskRunner
preprocessor = UnifiedIOPreprocessor.from_pretrained("allenai/uio2-preprocessor", tokenizer=tokenizer_path)
model = UnifiedIOModel.from_pretrained("allenai/uio2-large")
runner = TaskRunner(model, preprocessor)
image = runner.image_generation("a cat")
wavform = runner.audio_generation("dogs barking")

I am getting the error as

AttributeError: 'Decoder' object has no attribute 'generation_config'

Could you please let me know how to provide the generation_config?

Also, I tried the training code provided here as

from torch.nn import functional as F
from uio2.preprocessing import build_batch
preprocessed_example = preprocessor(
text_inputs="What is 1+1?", text_targets="2", target_modality="text")
batch = build_batch([preprocessed_example], device=model.device)
out = model(batch)
total_loss = 0
for modality, (logits, targets, mask) in out.items():
    losses = F.cross_entropy(
      logits.view(-1, logits.shape[-1]), targets.view(-1).to(torch.long), reduction="none")
    total_loss += (losses.reshape(logits.shape[:2])*mask)/mask.sum()
print(total_loss)

I am getting the IndexError while calculating out = model(batch) as:

/unified-io-2.pytorch/uio2/seq_features.py:160), in concat_sequences(seqs)
seq_lens = [x.seq_len for x in seqs]
out = {}
for k in dataclasses.fields(seqs[0]):
   k = k.name
   args = [expand_scalar(getattr(seq, k), seq.seq_len) for seq in seqs]
IndexError: list index out of range

When I tried to call generation as:

from uio2.preprocessing import build_batch 
preprocessed_example = preprocessor(text_inputs="What color is the sky?", target_modality="text")
batch = build_batch([preprocessed_example], device=model.device)
tokens = model.generate(batch, modality="text", max_new_tokens=128)

I am getting the error (even though I provided modality as text) as:

unified-io-2.pytorch/uio2/model.py:551 
if modality not in self.target_embedders:
    ValueError(f"No target encoder for {modality}")

allenai / unified-io-2.pytorch

Not able to run the Taskrunner or training code #7