LLM Aligner Bugs

[abhishekreddyallu]

LLM aligner workflows have generation and model-loading inconsistencies

1. Encoder-decoder output is sliced as decoder-only output

Location: ontoaligner/aligner/llm/llm.py:259, ontoaligner/aligner/llm/llm.py:285

Current behavior:

generated_ids = sequence_ids["sequences"][:, tokenized_input_data.input_ids.shape[-1]:]

This assumes the generated sequence contains:

prompt tokens + generated tokens

That is correct for decoder-only models, but not for encoder-decoder models such as Flan-T5.

Encoder-decoder models return only the generated target sequence. Therefore, slicing by the input prompt length removes the actual answer.

Observed before fix:

generate() received max_new_tokens: 5
Generated IDs passed to tokenizer: [[]]
Decoded output: ['']
AssertionError: BUG: encoder-decoder output was sliced like decoder-only output. T5-style models return only generated tokens, so prompt-length slicing removed the answer.

Fix:

Override encoder-decoder generation so it does not remove the prompt prefix:

class EncoderDecoderLLMArch(BaseLLMArch):
            generate_for_one_input()
                            ...
                       generated_ids = sequence_ids["sequences"]
            
            generate_for_multiple_input()
                            ...
                       generated_ids = sequence_ids["sequences"]

For encoder-decoder models, the generated sequence should be decoded directly.

Verified after fix:

generate() received max_new_tokens: 5
Generated IDs passed to tokenizer: [[101]]
Decoded output: ['yes']
PASS: encoder-decoder generation keeps the generated sequence.

---

2. Non-CPU Flan loading references undefined self.path

Location: ontoaligner/aligner/llm/llm.py:234

Current code:

            self.model = self.model.from_pretrained(self.path, load_in_8bit=True, device_map=self.kwargs['device_map'])

This branch uses self.path, but self.path is not defined on the Flan/T5 wrapper.

The method receives the model path as an argument:

def load_model(self, path: str):

So the supplied path should be forwarded directly.

Observed before fix:

AttributeError: 'FlanT5LEncoderDecoderLM' object has no attribute 'path'

This happens when loading Flan with a non-CPU device configuration.

Fix:

            self.model = self.model.from_pretrained(path, load_in_8bit=True, device_map=self.kwargs['device_map'])

)

This forwards the user-supplied model path correctly.

Verified after fix:

from_pretrained received path: 'qa-flan'
PASS: non-CPU Flan loading forwarded the supplied path correctly.

---

3. Batched local generation ignores num_beams

Location: ontoaligner/aligner/llm/llm.py:274-284

Current batched generation forwards:

max_new_tokens=self.kwargs["max_new_tokens"],
temperature=self.kwargs["temperature"],
top_p=self.kwargs["top_p"],
pad_token_id=self.tokenizer.eos_token_id,
...

But it does not forward:

num_beams=self.kwargs["num_beams"]

The one-input generation path forwards num_beams, but the multi-input generation path omits it. This causes inconsistent behavior depending on whether the user generates one prompt or a batch of prompts.

Fix:

    num_beams=self.kwargs["num_beams"]