Heuristic: does the model expect input_ids (integer tokens)?
(model: nn.Module)
| 311 | # --------------------------------------------------------------------------- |
| 312 | |
| 313 | def _is_language_model(model: nn.Module) -> bool: |
| 314 | """Heuristic: does the model expect input_ids (integer tokens)?""" |
| 315 | # Check common class names |
| 316 | cls_name = type(model).__name__.lower() |
| 317 | lm_indicators = [ |
| 318 | "causal", "lm", "gpt", "llama", "bert", "t5", "opt", "falcon", |
| 319 | "mistral", "gemma", "phi", "qwen", "codegen", "bloom", "mpt", |
| 320 | "seq2seq", |
| 321 | ] |
| 322 | if any(ind in cls_name for ind in lm_indicators): |
| 323 | return True |
| 324 | |
| 325 | # Check if model has an embedding layer as first child |
| 326 | for name, child in model.named_children(): |
| 327 | child_name = type(child).__name__.lower() |
| 328 | if "embed" in name.lower() or "embedding" in child_name: |
| 329 | return True |
| 330 | break # only check first child |
| 331 | |
| 332 | # Check forward signature for 'input_ids' |
| 333 | try: |
| 334 | sig = inspect.signature(model.forward) |
| 335 | if "input_ids" in sig.parameters: |
| 336 | return True |
| 337 | except (ValueError, TypeError): |
| 338 | pass |
| 339 | |
| 340 | return False |
| 341 | |
| 342 | |
| 343 | def generate_input( |