Fix return type annotation for TextClassificationModel.forward

[Copilot]

The forward method returns either torch.Tensor or dict[str, torch.Tensor] depending on return_label_attention_matrix, but was annotated as only returning torch.Tensor.

Changes:

• Updated return type to Union[torch.Tensor, dict[str, torch.Tensor]]
• Enhanced docstring to document both return paths with dict key structure
• Added validation: raises ValueError when return_label_attention_matrix=True without a configured text_embedder

Example:

# Now correctly typed for both cases
logits = model(input_ids, attention_mask, categorical_vars)  # torch.Tensor
result = model(input_ids, attention_mask, categorical_vars, return_label_attention_matrix=True)
# result: dict with keys "logits" and "label_attention_matrix"

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.