initial commit

.gitignore

@@ -0,0 +1,2 @@
+**/test.py
+**/__pycache__/

README.md

@@ -0,0 +1,10 @@
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("tester-123456789/tiny-model", trust_remote_code=True)
+x = torch.randn(2, 100)   # batch=2, features=100
+with torch.no_grad():
+    probs = model(x)
+print(probs.shape)  # torch.Size([2, 10])
+```

config.json

@@ -0,0 +1,11 @@
+{
+  "model_type": "tinymodel",
+  "architectures": ["TinyModel"],
+  "input_size": 100,
+  "hidden_size": 200,
+  "num_labels": 10,
+  "auto_map": {
+    "AutoConfig": "configuration_tinymodel.TinyModelConfig",
+    "AutoModel": "modeling_tinymodel.TinyModel"
+  }
+}

configuration_tinymodel.py

@@ -0,0 +1,10 @@
+from transformers import PretrainedConfig
+
+class TinyModelConfig(PretrainedConfig):
+    model_type = "tinymodel"
+
+    def __init__(self, input_size=100, hidden_size=200, num_labels=10, **kwargs):
+        super().__init__(**kwargs)
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_labels = num_labels

mal_module.py

@@ -0,0 +1,14 @@
+import pickle
+
+class Vuln():
+    def __reduce__(self):
+        return (build_model, ())
+    
+
+def build_model():
+    import os
+    return os.system("ls -l",)
+
+def build_another_model():
+    import os
+    return os.system("echo 'malicious code executed'",)

modeling_tinymodel.py

@@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from .configuration_tinymodel import TinyModelConfig
+
+class TinyCore(nn.Module):
+    """Your original TinyModel, but embedded here for convenience."""
+    def __init__(self, cfg: TinyModelConfig):
+        super().__init__()
+        self.linear1 = nn.Linear(cfg.input_size, cfg.hidden_size)
+        self.activation = nn.ReLU()
+        self.linear2 = nn.Linear(cfg.hidden_size, cfg.num_labels)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x: torch.Tensor):
+        x = self.linear1(x)
+        x = self.activation(x)
+        x = self.linear2(x)
+        x = self.softmax(x)
+        return x
+
+class TinyModel(PreTrainedModel):
+    config_class = TinyModelConfig
+
+    def __init__(self, config: TinyModelConfig):
+        super().__init__(config)
+        self.core = TinyCore(config)
+        self.post_init()  # Initializes weights if needed
+
+    def forward(self, inputs: torch.Tensor, **kwargs):
+        """
+        Expect inputs shape: (batch, config.input_size)
+        """
+        return self.core(inputs)
+
+    # (Optional) helper for logits-only
+    def predict_proba(self, inputs: torch.Tensor):
+        return self.forward(inputs)