trl-internal-testing/tiny-RemoteForCausalLM

1# Copyright 2020-2026 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import torch.nn as nn16from transformers import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaPreTrainedModel17 18from .configuration_remote import RemoteConfig19 20 21class RemoteModel(LlamaModel):22    config_class = RemoteConfig23 24 25class RemoteForCausalLM(LlamaForCausalLM):26    config_class = RemoteConfig27 28 29class RemoteForSequenceClassification(LlamaForSequenceClassification):30    config_class = RemoteConfig31 32    def __init__(self, config):33        # The parent's MRO calls `AutoModel.from_config(config)`, which would re-trigger the34        # trust_remote_code prompt for `RemoteConfig`. Wire `RemoteModel` in directly instead.35        LlamaPreTrainedModel.__init__(self, config)36        self.num_labels = config.num_labels37        self.model = RemoteModel(config)38        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)39        self.post_init()40