Fix .generate() support: add GenerationMixin + generation_config, document KV-cache limitation

#2

Summary

Calling .generate() on this model currently fails immediately with:

AttributeError: 'NanochatForCausalLM' object has no attribute 'generate'

Root cause: NanochatForCausalLM defines prepare_inputs_for_generation
(a strong "this model supports generation" signal to transformers) but
does not inherit from GenerationMixin. As of transformers>=4.50,
PreTrainedModel no longer automatically provides .generate() — this is
a breaking change that's been deprecated-warned for a few versions and is
now enforced. Your own trust_remote_code=True loading path already emits
this exact warning every time the model loads, so this isn't a surprise
bug — it's the deprecation landing.

Reproduction

```python
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

repo = "himalaya-ai/himalayagpt-0.5b-it"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
repo, trust_remote_code=True, torch_dtype=torch.float32, device_map="auto"
)
ids = tok("नेपालको राजधानी ", return_tensors="pt").input_ids.to(model.device)
out = model.generate(ids, max_new_tokens=64)

AttributeError: 'NanochatForCausalLM' object has no attribute 'generate'

```

Fix in this PR

  1. Add GenerationMixin to the class bases so .generate() exists:
    ```python
    class NanochatForCausalLM(PreTrainedModel, GenerationMixin):
    ```

  2. Set a default generation_config in __init__, since nothing
    currently provides one and generate() reads from it immediately
    (self.generation_config.max_length, etc.):
    ```python
    self.generation_config = GenerationConfig.from_model_config(config)
    ```

That's the minimal fix to make .generate() callable at all.

Known limitation this PR does NOT fix (flagging for visibility)

prepare_inputs_for_generation currently always returns the full
input_ids regardless of past_key_values, and forward() hardcodes
past_key_values=None in its return — meaning there is no KV-cache.
Generation will work correctly after this PR, but every new token
reprocesses the entire sequence from scratch (O(n²) instead of O(n)). For
a 0.5B model at short context this is usable but slow; for longer
generations it'll be noticeably worse than a cached implementation.

I'd suggest treating that as a separate, bigger follow-up PR (it requires
threading cache state through NanochatAttention.forward, which doesn't
currently accept or return any cache tensors) rather than blocking this
fix on it. Happy to take a pass at that separately if useful — wanted to
keep this PR small and easy to review/merge.

Also worth noting: forward() accepts an attention_mask parameter but
never uses it — batched/padded generation will silently ignore padding.
Only affects multi-sequence batched calls, not single-prompt generation.

Testing

After this change, confirmed locally:

  • model.generate(ids, max_new_tokens=64) runs without raising
  • Output is well-formed token ids, decodes via tok.decode(...) normally
  • forward() / loss computation (training path) is untouched — only the
    class declaration and __init__ changed
Himalaya AI Labs org

LGTM. Verified that the PR fixes direct .generate() support by adding GenerationMixin and a default GenerationConfig, without touching the model forward path, attention logic, rotary code, or weights.

HimalayaGPT changed pull request status to merged

Sign up or log in to comment