File size: 579 Bytes
d6e983a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# modeling.py
from transformers import PreTrainedModel
import torch.nn as nn

class SimpleWarningConfig:
    def __init__(self, hidden_size=8):
        self.hidden_size = hidden_size

class SimpleModel(PreTrainedModel):
    config_class = SimpleWarningConfig

    def __init__(self, config):
        super().__init__(config)
        print("⚠️ WARNING: You are loading a model with trust_remote_code=True! Be careful with remote code execution!")
        self.linear = nn.Linear(config.hidden_size, config.hidden_size)

    def forward(self, x):
        return self.linear(x)