Wrap ForwardContext around full model forward

[calpt]

This PR adapts the ForwardContext to be applied to the full model (including head) forward pass. The original base model forward wrapper is now moved to wrap_base to make sure no second ForwardContext is created for a single forward pass.

This enables passing custom args that are defined in the ForwardContext definition to the top-level model call, as discussed in #783, e.g.:

model = AutoModelForCausalLM.from_pretrained(model_name)
adapters.init(model)
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer(["This is a test text"], return_tensors="pt")

# Registers new forward args globally
ForwardContext.context_args += ["task_ids"]

# New the new arg name can be used w/o modifying the model's forward method
output = model(**inputs, task_ids=["id_0", "id_1"])

In the example above, the forward context will automatically add the passed context args as attributes, ie. they can be accessed within the foward pass like this:

task_ids = ForwardContext.get_context().task_ids

[FrLdy]

Hello @calpt,
Thanks for the PR !
Just tested it out and it works great !
This definitely makes passing custom args much easier, no need to redefine a new forward for each specific hf model.
I'll rebase my fork once the changes are merged and propose my additions for multi-task support.

[FrLdy]

Hi,
I tried to include a new ForwardContext in my test for multi-task composition:
https://github.com/FrLdy/adapters/blob/6f63c4df35d8deed973f4c791f6b779f8ba4f668/tests/test_misc/test_adapter_composition.py#L177

However, if I don't check whether the new ForwardContext argument (task_ids) is present in context_args, the test runs fine when each TestCase class is executed independently. But when they are run together, only the first one passes, while the subsequent ones fail.

Maybe use a set to store ForwardContext.context_args ?